The forward entry with one level of blocks expanded. Every edge carries the tensor type the compiler inferred at that point, in the model's own generics.

Entries

Entry	Signature
forward	forward<B: Dim, S: Dim>(tokens: Tensor[B, S; i32]) -> Tensor[B, S, Vocab; T]
next_token	next_token<B: Dim, S: Dim>(tokens: Tensor[B, S; i32]) -> Tensor[B, Vocab; T]
decode	decode(token: Tensor[Batch, 1; i32], pos: i32) -> Tensor[Batch, Vocab; T]
prefill	prefill<S: Dim>(tokens: Tensor[Batch, S; i32], pos: i32) -> Tensor[Batch, Vocab; T]
decode_rows	decode_rows(tokens: Tensor[Batch, 1; i32], positions: Tensor[Batch; i32]) -> Tensor[Batch, Vocab; T]
prefill_slots	prefill_slots<M: Dim, S: Dim>(tokens: Tensor[M, S; i32], slots: Tensor[M; i32], lengths: Tensor[M; i32]) -> Tensor[M, Vocab; T]
loss_packed	loss_packed<P: Dim>(tokens: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32], targets: Tensor[P; i64], weights: Tensor[P; f32]) -> f32
log_probs_packed	log_probs_packed<P: Dim>(tokens: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32], targets: Tensor[P; i64]) -> Tensor[P; f32]
hidden_packed	hidden_packed<P: Dim>(tokens: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32]) -> Tensor[P, H; T]
prefill_packed	prefill_packed<P: Dim>(tokens: Tensor[P; i32], rows: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32], last: Tensor[Batch; i32]) -> Tensor[Batch, Vocab; T]
step_packed	step_packed<P: Dim>(tokens: Tensor[P; i32], rows: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32], last: Tensor[Batch; i32], step_tokens: Tensor[Batch, 1; i32], step_positions: Tensor[Batch; i32]) -> Tensor[2 * Batch, Vocab; T]
generate	generate<Steps: Dim>(token: Tensor[Batch, 1; i32], pos: i32) -> Tensor[Batch, Steps; i32]
sample	sample<Steps: Dim>(token: Tensor[Batch, 1; i32], pos: i32, key: Tensor[2; i64], temperature: f32) -> Tensor[Batch, Steps; i32]
generate_until	generate_until<MaxNew: Dim>(token: Tensor[Batch, 1; i32], pos: i32, eos: i32) -> (Tensor[Batch, MaxNew; i32], i32)

Generics

The root block's generics as this checkpoint binds them.

Vocab	32064
H	3072
Heads	32
Inner	8192
Layers	32
Batch	1
MaxSeq	4096
T	bf16

Blocks

Every block of the program with its members and functions, as linnet inspect prints them.

Attention

phi3.attention::Attention<H: Dim, Heads: Dim, Batch: Dim, MaxSeq: Dim, T: Float>
  sub qkv: Linear<H, 3 * H, T>
  sub o_proj: Linear<H, H, T>
  state cache_k: Tensor[Batch, Heads, MaxSeq, H / Heads; T]
  state cache_v: Tensor[Batch, Heads, MaxSeq, H / Heads; T]
  pub fn forward<B: Dim, S: Dim>(x: Tensor[B, S, H; T]) -> Tensor[B, S, H; T]
  pub fn decode(x: Tensor[Batch, 1, H; T], pos: i32) -> Tensor[Batch, 1, H; T]
  pub fn prefill<S: Dim>(x: Tensor[Batch, S, H; T], pos: i32) -> Tensor[Batch, S, H; T]
  pub fn decode_rows(x: Tensor[Batch, 1, H; T], positions: Tensor[Batch; i32]) -> Tensor[Batch, 1, H; T]
  pub fn prefill_slots<M: Dim, S: Dim>(x: Tensor[M, S, H; T], slots: Tensor[M; i32]) -> Tensor[M, S, H; T]
  pub fn prefill_packed<P: Dim>(x: Tensor[1, P, H; T], rows: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32]) -> Tensor[1, P, H; T]
  pub fn packed<P: Dim>(x: Tensor[1, P, H; T], positions: Tensor[P; i32], segments: Tensor[P; i32]) -> Tensor[1, P, H; T]
  pub fn step_packed<P: Dim>(x: Tensor[1, Batch + P, H; T], rows: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32], step_positions: Tensor[Batch; i32]) -> Tensor[1, Batch + P, H; T]

DecoderLayer

phi3::DecoderLayer<H: Dim, Heads: Dim, Inner: Dim, Batch: Dim, MaxSeq: Dim, T: Float>
  sub attention_norm: RmsNorm<H, T>
  sub attention: Attention<H, Heads, Batch, MaxSeq, T>
  sub mlp_norm: RmsNorm<H, T>
  sub mlp: FusedSwiGlu<H, Inner, T>
  pub fn forward<B: Dim, S: Dim>(x: Tensor[B, S, H; T]) -> Tensor[B, S, H; T]
  pub fn decode(x: Tensor[Batch, 1, H; T], pos: i32) -> Tensor[Batch, 1, H; T]
  pub fn prefill<S: Dim>(x: Tensor[Batch, S, H; T], pos: i32) -> Tensor[Batch, S, H; T]
  pub fn decode_rows(x: Tensor[Batch, 1, H; T], positions: Tensor[Batch; i32]) -> Tensor[Batch, 1, H; T]
  pub fn prefill_slots<M: Dim, S: Dim>(x: Tensor[M, S, H; T], slots: Tensor[M; i32]) -> Tensor[M, S, H; T]
  pub fn prefill_packed<P: Dim>(x: Tensor[1, P, H; T], rows: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32]) -> Tensor[1, P, H; T]
  pub fn packed<P: Dim>(x: Tensor[1, P, H; T], positions: Tensor[P; i32], segments: Tensor[P; i32]) -> Tensor[1, P, H; T]
  pub fn step_packed<P: Dim>(x: Tensor[1, Batch + P, H; T], rows: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32], step_positions: Tensor[Batch; i32]) -> Tensor[1, Batch + P, H; T]

Embedding

std.nn.embedding::Embedding<Vocab: Dim, H: Dim, T: Float = bf16>
  param weight: Tensor[Vocab, H; T]
  pub fn forward<*S: Shape>(ids: Tensor[*S; i32]) -> Tensor[*S, H; T]

FusedSwiGlu

phi3.mlp::FusedSwiGlu<H: Dim, Inner: Dim, T: Float = bf16>
  sub gate_up: Linear<H, 2 * Inner, T>
  sub down: Linear<Inner, H, T>
  pub fn forward<*S: Shape>(x: Tensor[*S, H; T]) -> Tensor[*S, H; T]

Linear

std.nn.linear::Linear<In: Dim, Out: Dim, T: Float = bf16>
  param weight: Tensor[Out, In; T]
  param bias: Tensor[Out; T]?
  pub fn forward<*S: Shape>(x: Tensor[*S, In; T]) -> Tensor[*S, Out; T]

Model

phi3::Model<Vocab: Dim, H: Dim, Heads: Dim, Inner: Dim, Layers: Dim, Batch: Dim, MaxSeq: Dim, T: Float = bf16>
  sub embedding: Embedding<Vocab, H, T>
  sub layers: [DecoderLayer<H, Heads, Inner, Batch, MaxSeq, T>; Layers]
  sub norm: RmsNorm<H, T>
  sub lm_head: Linear<H, Vocab, T>
  pub entry forward<B: Dim, S: Dim>(tokens: Tensor[B, S; i32]) -> Tensor[B, S, Vocab; T]
  pub entry next_token<B: Dim, S: Dim>(tokens: Tensor[B, S; i32]) -> Tensor[B, Vocab; T]
  pub entry decode(token: Tensor[Batch, 1; i32], pos: i32) -> Tensor[Batch, Vocab; T]
  pub entry prefill<S: Dim>(tokens: Tensor[Batch, S; i32], pos: i32) -> Tensor[Batch, Vocab; T]
  pub entry decode_rows(tokens: Tensor[Batch, 1; i32], positions: Tensor[Batch; i32]) -> Tensor[Batch, Vocab; T]
  pub entry prefill_slots<M: Dim, S: Dim>(tokens: Tensor[M, S; i32], slots: Tensor[M; i32], lengths: Tensor[M; i32]) -> Tensor[M, Vocab; T]
  pub entry loss_packed<P: Dim>(tokens: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32], targets: Tensor[P; i64], weights: Tensor[P; f32]) -> f32
  pub entry log_probs_packed<P: Dim>(tokens: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32], targets: Tensor[P; i64]) -> Tensor[P; f32]
  pub entry hidden_packed<P: Dim>(tokens: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32]) -> Tensor[P, H; T]
  fn packed_states<P: Dim>(tokens: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32]) -> Tensor[P, H; T]
  pub entry prefill_packed<P: Dim>(tokens: Tensor[P; i32], rows: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32], last: Tensor[Batch; i32]) -> Tensor[Batch, Vocab; T]
  pub entry step_packed<P: Dim>(tokens: Tensor[P; i32], rows: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32], last: Tensor[Batch; i32], step_tokens: Tensor[Batch, 1; i32], step_positions: Tensor[Batch; i32]) -> Tensor[2 * Batch, Vocab; T]
  pub entry generate<Steps: Dim>(token: Tensor[Batch, 1; i32], pos: i32) -> Tensor[Batch, Steps; i32]
  pub entry sample<Steps: Dim>(token: Tensor[Batch, 1; i32], pos: i32, key: Tensor[2; i64], temperature: f32) -> Tensor[Batch, Steps; i32]
  pub entry generate_until<MaxNew: Dim>(token: Tensor[Batch, 1; i32], pos: i32, eos: i32) -> (Tensor[Batch, MaxNew; i32], i32)
  fn step(token: Tensor[Batch, 1; i32], pos: i32) -> Tensor[Batch, Vocab; T]
  fn hidden<B: Dim, S: Dim>(tokens: Tensor[B, S; i32]) -> Tensor[B, S, H; T]
  fn hidden_states<B: Dim, S: Dim>(tokens: Tensor[B, S; i32]) -> Tensor[B, S, H; T]

RmsNorm

std.nn.norm::RmsNorm<H: Dim, T: Float = bf16>
  param weight: Tensor[H; T]
  pub fn forward<*S: Shape>(x: Tensor[*S, H; T]) -> Tensor[*S, H; T]