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>(image: Tensor[B, Channels, Height, Width; T]) -> Tensor[B, Classes; T]

Generics

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

Height	224
Width	224
Channels	3
Patch	16
D	768
Heads	12
Inner	3072
Layers	12
Classes	1000
T	f32

Blocks

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

Attention

vit::Attention<D: Dim, Heads: Dim, T: Float>
  sub query: Linear<D, D, T>
  sub key: Linear<D, D, T>
  sub value: Linear<D, D, T>
  sub out: Linear<D, D, T>
  pub fn forward<B: Dim, N: Dim>(x: Tensor[B, N, D; T]) -> Tensor[B, N, D; T]

EncoderLayer

vit::EncoderLayer<D: Dim, Heads: Dim, Inner: Dim, T: Float>
  sub norm_before: LayerNorm<D, T>
  sub attention: Attention<D, Heads, T>
  sub norm_after: LayerNorm<D, T>
  sub up: Linear<D, Inner, T>
  sub down: Linear<Inner, D, T>
  pub fn forward<B: Dim, N: Dim>(x: Tensor[B, N, D; T]) -> Tensor[B, N, D; T]

LayerNorm

vit::LayerNorm<D: Dim, T: Float>
  param weight: Tensor[D; T]
  param bias: Tensor[D; T]
  pub fn forward<*S: Shape>(x: Tensor[*S, D; T]) -> Tensor[*S, D; 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

vit::Model<Height: Dim, Width: Dim, Channels: Dim, Patch: Dim, D: Dim, Heads: Dim, Inner: Dim, Layers: Dim, Classes: Dim, T: Float = f32>
  param patch_weight: Tensor[D, Channels, Patch, Patch; T]
  param patch_bias: Tensor[D; T]
  param class_token: Tensor[1, 1, D; T]
  param positions: Tensor[1, 1 + (Height / Patch) * (Width / Patch), D; T]
  sub layers: [EncoderLayer<D, Heads, Inner, T>; Layers]
  sub norm: LayerNorm<D, T>
  sub head: Linear<D, Classes, T>
  pub entry forward<B: Dim>(image: Tensor[B, Channels, Height, Width; T]) -> Tensor[B, Classes; T]
  fn embed<B: Dim>(image: Tensor[B, Channels, Height, Width; T]) -> Tensor[B, (Height / Patch) * (Width / Patch), D; 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]