The decode 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
decode	decode<B: Dim>(latent: Tensor[B, 4, H, W; T]) -> Tensor[B, 3, 8 * H, 8 * W; T]

Generics

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

H	16
W	16
T	f32

Blocks

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

Conv1d

std.nn.conv::Conv1d<Cin: Dim, Cout: Dim, K: Dim, Stride: Dim, Pad: Dim, T: Float = f32>
  param weight: Tensor[Cout, Cin, K; T]
  param bias: Tensor[Cout; T]?
  pub fn forward<B: Dim, L: Dim>(x: Tensor[B, Cin, L; T]) -> Tensor[B, Cout, 1 + (-1 * K + L + 2 * Pad) / Stride; T]

Conv2d

std.nn.conv::Conv2d<Cin: Dim, Cout: Dim, K: Dim, Stride: Dim, Pad: Dim, T: Float = f32>
  param weight: Tensor[Cout, Cin, K, K; T]
  param bias: Tensor[Cout; T]?
  pub fn forward<B: Dim, H: Dim, W: Dim>(x: Tensor[B, Cin, H, W; T]) -> Tensor[B, Cout, 1 + (-1 * K + H + 2 * Pad) / Stride, 1 + (-1 * K + W + 2 * Pad) / Stride; T]

Conv2dRect

std.nn.conv::Conv2dRect<Cin: Dim, Cout: Dim, KH: Dim, KW: Dim, StrideH: Dim, StrideW: Dim, PadH: Dim, PadW: Dim, T: Float = f32>
  param weight: Tensor[Cout, Cin, KH, KW; T]
  param bias: Tensor[Cout; T]?
  pub fn forward<B: Dim, H: Dim, W: Dim>(x: Tensor[B, Cin, H, W; T]) -> Tensor[B, Cout, 1 + (-1 * KH + H + 2 * PadH) / StrideH, 1 + (-1 * KW + W + 2 * PadW) / StrideW; T]

GroupNorm

sd_vae::GroupNorm<C: Dim, T: Float>
  param weight: Tensor[C; T]
  param bias: Tensor[C; T]
  pub fn forward<B: Dim, H: Dim, W: Dim>(x: Tensor[B, C, H, W; T]) -> Tensor[B, C, H, W; 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]

MidBlock

sd_vae::MidBlock<C: Dim, T: Float>
  sub first: ResBlock<C, T>
  sub attention: SpatialAttention<C, T>
  sub second: ResBlock<C, T>
  pub fn forward<B: Dim, H: Dim, W: Dim>(x: Tensor[B, C, H, W; T]) -> Tensor[B, C, H, W; T]

Model

sd_vae::Model<H: Dim, W: Dim, T: Float = f32>
  sub post_quant_conv: Conv2d<4, 4, 1, 1, 0, T>
  sub conv_in: Conv2d<4, 512, 3, 1, 1, T>
  sub mid_block: MidBlock<512, T>
  sub up_block_0: UpBlock<512, T>
  sub up_block_1: UpBlock<512, T>
  sub up_block_2: UpBlockWiden<512, 256, T>
  sub up_block_3: UpBlockWidenFinal<256, 128, T>
  sub norm_out: GroupNorm<128, T>
  sub conv_out: Conv2d<128, 3, 3, 1, 1, T>
  pub entry decode<B: Dim>(latent: Tensor[B, 4, H, W; T]) -> Tensor[B, 3, 8 * H, 8 * W; T]

ResBlock

sd_vae::ResBlock<C: Dim, T: Float>
  sub norm1: GroupNorm<C, T>
  sub conv1: Conv2d<C, C, 3, 1, 1, T>
  sub norm2: GroupNorm<C, T>
  sub conv2: Conv2d<C, C, 3, 1, 1, T>
  pub fn forward<B: Dim, H: Dim, W: Dim>(x: Tensor[B, C, H, W; T]) -> Tensor[B, C, H, W; T]

ResBlockWiden

sd_vae::ResBlockWiden<Cin: Dim, Cout: Dim, T: Float>
  sub norm1: GroupNorm<Cin, T>
  sub conv1: Conv2d<Cin, Cout, 3, 1, 1, T>
  sub norm2: GroupNorm<Cout, T>
  sub conv2: Conv2d<Cout, Cout, 3, 1, 1, T>
  sub shortcut: Conv2d<Cin, Cout, 1, 1, 0, T>
  pub fn forward<B: Dim, H: Dim, W: Dim>(x: Tensor[B, Cin, H, W; T]) -> Tensor[B, Cout, H, W; 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]

SpatialAttention

sd_vae::SpatialAttention<C: Dim, T: Float>
  sub norm: GroupNorm<C, T>
  sub query: Linear<C, C, T>
  sub key: Linear<C, C, T>
  sub value: Linear<C, C, T>
  sub proj: Linear<C, C, T>
  pub fn forward<B: Dim, H: Dim, W: Dim>(x: Tensor[B, C, H, W; T]) -> Tensor[B, C, H, W; T]

UpBlock

sd_vae::UpBlock<C: Dim, T: Float>
  sub first: ResBlock<C, T>
  sub second: ResBlock<C, T>
  sub third: ResBlock<C, T>
  sub upsample: Conv2d<C, C, 3, 1, 1, T>
  pub fn forward<B: Dim, H: Dim, W: Dim>(x: Tensor[B, C, H, W; T]) -> Tensor[B, C, 2 * H, 2 * W; T]

UpBlockWiden

sd_vae::UpBlockWiden<Cin: Dim, Cout: Dim, T: Float>
  sub first: ResBlockWiden<Cin, Cout, T>
  sub second: ResBlock<Cout, T>
  sub third: ResBlock<Cout, T>
  sub upsample: Conv2d<Cout, Cout, 3, 1, 1, T>
  pub fn forward<B: Dim, H: Dim, W: Dim>(x: Tensor[B, Cin, H, W; T]) -> Tensor[B, Cout, 2 * H, 2 * W; T]

UpBlockWidenFinal

sd_vae::UpBlockWidenFinal<Cin: Dim, Cout: Dim, T: Float>
  sub first: ResBlockWiden<Cin, Cout, T>
  sub second: ResBlock<Cout, T>
  sub third: ResBlock<Cout, T>
  pub fn forward<B: Dim, H: Dim, W: Dim>(x: Tensor[B, Cin, H, W; T]) -> Tensor[B, Cout, H, W; T]