Upload model at step 45000

config.py

@@ -6,18 +6,2037 @@ Precision: float32
 """
 
 from triax.models.nn.condUNet import CondUNet2D
+import jax.numpy as jnp
+from flax import nnx
 
 # Model architecture
-# TODO: Fill in the actual initialization parameters from your training config
 model = CondUNet2D(
-    # Add your model parameters here
-    # Example:
-    # hidden_dim=256,
-    # num_layers=4,
-    # etc.
+    blocks_down=[TimestepEmbedSequential( # Param: 320 (1.3 KB)
+  layers=[Conv( # Param: 320 (1.3 KB)
+    bias=Param( # 32 (128 B)
+      value=Array(shape=(32,), dtype=dtype('float32'))
+    ),
+    bias_init=<function zeros at 0x7fb32b98c2c0>,
+    conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+    dtype=float32,
+    feature_group_count=1,
+    in_features=1,
+    input_dilation=1,
+    kernel=Param( # 288 (1.2 KB)
+      value=Array(shape=(3, 3, 1, 32), dtype=dtype('float32'))
+    ),
+    kernel_dilation=1,
+    kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+    kernel_shape=(3, 3, 1, 64),
+    kernel_size=(3, 3),
+    mask=None,
+    out_features=64,
+    padding=(1, 1),
+    param_dtype=float32,
+    precision=None,
+    promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+    strides=(1, 1),
+    use_bias=True
+  )]
+), TimestepEmbedSequential( # Param: 26,880 (107.5 KB), RngState: 2 (12 B), Total: 26,882 (107.5 KB)
+  layers=[ResnetBlock( # Param: 26,880 (107.5 KB), RngState: 2 (12 B), Total: 26,882 (107.5 KB)
+    activation=<PjitFunction of <function silu at 0x7fb32b8440e0>>,
+    conv1=Conv( # Param: 9,248 (37.0 KB)
+      bias=Param( # 32 (128 B)
+        value=Array(shape=(32,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=64,
+      input_dilation=1,
+      kernel=Param( # 9,216 (36.9 KB)
+        value=Array(shape=(3, 3, 32, 32), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(3, 3, 64, 64),
+      kernel_size=(3, 3),
+      mask=None,
+      out_features=64,
+      padding=(1, 1),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    conv2=Conv( # Param: 9,248 (37.0 KB)
+      bias=Param( # 32 (128 B)
+        value=Array(shape=(32,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=64,
+      input_dilation=1,
+      kernel=Param( # 9,216 (36.9 KB)
+        value=Array(shape=(3, 3, 32, 32), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(3, 3, 64, 64),
+      kernel_size=(3, 3),
+      mask=None,
+      out_features=64,
+      padding=(1, 1),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    dropout=Dropout( # RngState: 2 (12 B)
+      broadcast_dims=(),
+      deterministic=False,
+      rate=0.1,
+      rng_collection='dropout',
+      rngs=RngStream( # RngState: 2 (12 B)
+        count=RngCount( # 1 (4 B)
+          value=Array(0, dtype=uint32),
+          tag='default'
+        ),
+        key=RngKey( # 1 (8 B)
+          value=Array((), dtype=key<fry>) overlaying:
+          [2585633080 2083471411],
+          tag='default'
+        ),
+        tag='default'
+      )
+    ),
+    embedding_dim=256,
+    in_channels=64,
+    norm1=GroupNorm( # Param: 64 (256 B)
+      axis_index_groups=None,
+      axis_name=None,
+      bias=Param( # 32 (128 B)
+        value=Array(shape=(32,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dtype=float32,
+      epsilon=1e-06,
+      feature_axis=-1,
+      group_size=2,
+      num_groups=32,
+      param_dtype=float32,
+      reduction_axes=None,
+      scale=Param( # 32 (128 B)
+        value=Array(shape=(32,), dtype=dtype('float32'))
+      ),
+      scale_init=<function ones at 0x7fb32b86e520>,
+      use_bias=True,
+      use_fast_variance=True,
+      use_scale=True
+    ),
+    norm2=GroupNorm( # Param: 64 (256 B)
+      axis_index_groups=None,
+      axis_name=None,
+      bias=Param( # 32 (128 B)
+        value=Array(shape=(32,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dtype=float32,
+      epsilon=1e-06,
+      feature_axis=-1,
+      group_size=2,
+      num_groups=32,
+      param_dtype=float32,
+      reduction_axes=None,
+      scale=Param( # 32 (128 B)
+        value=Array(shape=(32,), dtype=dtype('float32'))
+      ),
+      scale_init=<function ones at 0x7fb32b86e520>,
+      use_bias=True,
+      use_fast_variance=True,
+      use_scale=True
+    ),
+    out_channels=64,
+    time_mlp=Linear( # Param: 8,256 (33.0 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dot_general=<function dot_general at 0x7fb32c21a3e0>,
+      dtype=float32,
+      in_features=256,
+      kernel=Param( # 8,192 (32.8 KB)
+        value=Array(shape=(128, 64), dtype=dtype('float32'))
+      ),
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      out_features=128,
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      use_bias=True
+    )
+  )]
+), TimestepEmbedSequential( # Param: 9,248 (37.0 KB)
+  layers=[Downsample( # Param: 9,248 (37.0 KB)
+    conv=Conv( # Param: 9,248 (37.0 KB)
+      bias=Param( # 32 (128 B)
+        value=Array(shape=(32,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=64,
+      input_dilation=1,
+      kernel=Param( # 9,216 (36.9 KB)
+        value=Array(shape=(3, 3, 32, 32), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(3, 3, 64, 64),
+      kernel_size=(3, 3),
+      mask=None,
+      out_features=64,
+      padding=(0, 0),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(2, 2),
+      use_bias=True
+    ),
+    method='conv'
+  )]
+), TimestepEmbedSequential( # Param: 91,008 (364.0 KB), RngState: 2 (12 B), Total: 91,010 (364.0 KB)
+  layers=[ResnetBlock( # Param: 74,240 (297.0 KB), RngState: 2 (12 B), Total: 74,242 (297.0 KB)
+    activation=<PjitFunction of <function silu at 0x7fb32b8440e0>>,
+    conv1=Conv( # Param: 18,496 (74.0 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=64,
+      input_dilation=1,
+      kernel=Param( # 18,432 (73.7 KB)
+        value=Array(shape=(3, 3, 32, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(3, 3, 64, 128),
+      kernel_size=(3, 3),
+      mask=None,
+      out_features=128,
+      padding=(1, 1),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    conv2=Conv( # Param: 36,928 (147.7 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 36,864 (147.5 KB)
+        value=Array(shape=(3, 3, 64, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(3, 3, 128, 128),
+      kernel_size=(3, 3),
+      mask=None,
+      out_features=128,
+      padding=(1, 1),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    dropout=Dropout( # RngState: 2 (12 B)
+      broadcast_dims=(),
+      deterministic=False,
+      rate=0.1,
+      rng_collection='dropout',
+      rngs=RngStream( # RngState: 2 (12 B)
+        count=RngCount( # 1 (4 B)
+          value=Array(0, dtype=uint32),
+          tag='default'
+        ),
+        key=RngKey( # 1 (8 B)
+          value=Array((), dtype=key<fry>) overlaying:
+          [2656139193 2766658851],
+          tag='default'
+        ),
+        tag='default'
+      )
+    ),
+    embedding_dim=256,
+    in_channels=64,
+    nin_shortcut=Conv( # Param: 2,112 (8.4 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=64,
+      input_dilation=1,
+      kernel=Param( # 2,048 (8.2 KB)
+        value=Array(shape=(1, 1, 32, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(1, 1, 64, 128),
+      kernel_size=(1, 1),
+      mask=None,
+      out_features=128,
+      padding=(0, 0),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    norm1=GroupNorm( # Param: 64 (256 B)
+      axis_index_groups=None,
+      axis_name=None,
+      bias=Param( # 32 (128 B)
+        value=Array(shape=(32,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dtype=float32,
+      epsilon=1e-06,
+      feature_axis=-1,
+      group_size=2,
+      num_groups=32,
+      param_dtype=float32,
+      reduction_axes=None,
+      scale=Param( # 32 (128 B)
+        value=Array(shape=(32,), dtype=dtype('float32'))
+      ),
+      scale_init=<function ones at 0x7fb32b86e520>,
+      use_bias=True,
+      use_fast_variance=True,
+      use_scale=True
+    ),
+    norm2=GroupNorm( # Param: 128 (512 B)
+      axis_index_groups=None,
+      axis_name=None,
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dtype=float32,
+      epsilon=1e-06,
+      feature_axis=-1,
+      group_size=4,
+      num_groups=32,
+      param_dtype=float32,
+      reduction_axes=None,
+      scale=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      scale_init=<function ones at 0x7fb32b86e520>,
+      use_bias=True,
+      use_fast_variance=True,
+      use_scale=True
+    ),
+    out_channels=128,
+    time_mlp=Linear( # Param: 16,512 (66.0 KB)
+      bias=Param( # 128 (512 B)
+        value=Array(shape=(128,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dot_general=<function dot_general at 0x7fb32c21a3e0>,
+      dtype=float32,
+      in_features=256,
+      kernel=Param( # 16,384 (65.5 KB)
+        value=Array(shape=(128, 128), dtype=dtype('float32'))
+      ),
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      out_features=256,
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      use_bias=True
+    )
+  ), AttnBlock( # Param: 16,768 (67.1 KB)
+    dtype=float32,
+    head_dim=32,
+    k=Conv( # Param: 4,160 (16.6 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 4,096 (16.4 KB)
+        value=Array(shape=(1, 1, 64, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(1, 1, 128, 128),
+      kernel_size=(1, 1),
+      mask=None,
+      out_features=128,
+      padding='SAME',
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=1,
+      use_bias=True
+    ),
+    norm=GroupNorm( # Param: 128 (512 B)
+      axis_index_groups=None,
+      axis_name=None,
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dtype=float32,
+      epsilon=1e-06,
+      feature_axis=-1,
+      group_size=4,
+      num_groups=32,
+      param_dtype=float32,
+      reduction_axes=None,
+      scale=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      scale_init=<function ones at 0x7fb32b86e520>,
+      use_bias=True,
+      use_fast_variance=True,
+      use_scale=True
+    ),
+    num_heads=4,
+    proj_out=Conv( # Param: 4,160 (16.6 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 4,096 (16.4 KB)
+        value=Array(shape=(1, 1, 64, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(1, 1, 128, 128),
+      kernel_size=(1, 1),
+      mask=None,
+      out_features=128,
+      padding='SAME',
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=1,
+      use_bias=True
+    ),
+    q=Conv( # Param: 4,160 (16.6 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 4,096 (16.4 KB)
+        value=Array(shape=(1, 1, 64, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(1, 1, 128, 128),
+      kernel_size=(1, 1),
+      mask=None,
+      out_features=128,
+      padding='SAME',
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=1,
+      use_bias=True
+    ),
+    v=Conv( # Param: 4,160 (16.6 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 4,096 (16.4 KB)
+        value=Array(shape=(1, 1, 64, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(1, 1, 128, 128),
+      kernel_size=(1, 1),
+      mask=None,
+      out_features=128,
+      padding='SAME',
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=1,
+      use_bias=True
+    )
+  )]
+), TimestepEmbedSequential( # Param: 36,928 (147.7 KB)
+  layers=[Downsample( # Param: 36,928 (147.7 KB)
+    conv=Conv( # Param: 36,928 (147.7 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 36,864 (147.5 KB)
+        value=Array(shape=(3, 3, 64, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(3, 3, 128, 128),
+      kernel_size=(3, 3),
+      mask=None,
+      out_features=128,
+      padding=(0, 0),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(2, 2),
+      use_bias=True
+    ),
+    method='conv'
+  )]
+), TimestepEmbedSequential( # Param: 90,624 (362.5 KB), RngState: 2 (12 B), Total: 90,626 (362.5 KB)
+  layers=[ResnetBlock( # Param: 90,624 (362.5 KB), RngState: 2 (12 B), Total: 90,626 (362.5 KB)
+    activation=<PjitFunction of <function silu at 0x7fb32b8440e0>>,
+    conv1=Conv( # Param: 36,928 (147.7 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 36,864 (147.5 KB)
+        value=Array(shape=(3, 3, 64, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(3, 3, 128, 128),
+      kernel_size=(3, 3),
+      mask=None,
+      out_features=128,
+      padding=(1, 1),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    conv2=Conv( # Param: 36,928 (147.7 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 36,864 (147.5 KB)
+        value=Array(shape=(3, 3, 64, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(3, 3, 128, 128),
+      kernel_size=(3, 3),
+      mask=None,
+      out_features=128,
+      padding=(1, 1),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    dropout=Dropout( # RngState: 2 (12 B)
+      broadcast_dims=(),
+      deterministic=False,
+      rate=0.1,
+      rng_collection='dropout',
+      rngs=RngStream( # RngState: 2 (12 B)
+        count=RngCount( # 1 (4 B)
+          value=Array(0, dtype=uint32),
+          tag='default'
+        ),
+        key=RngKey( # 1 (8 B)
+          value=Array((), dtype=key<fry>) overlaying:
+          [ 692128146 1043829861],
+          tag='default'
+        ),
+        tag='default'
+      )
+    ),
+    embedding_dim=256,
+    in_channels=128,
+    norm1=GroupNorm( # Param: 128 (512 B)
+      axis_index_groups=None,
+      axis_name=None,
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dtype=float32,
+      epsilon=1e-06,
+      feature_axis=-1,
+      group_size=4,
+      num_groups=32,
+      param_dtype=float32,
+      reduction_axes=None,
+      scale=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      scale_init=<function ones at 0x7fb32b86e520>,
+      use_bias=True,
+      use_fast_variance=True,
+      use_scale=True
+    ),
+    norm2=GroupNorm( # Param: 128 (512 B)
+      axis_index_groups=None,
+      axis_name=None,
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dtype=float32,
+      epsilon=1e-06,
+      feature_axis=-1,
+      group_size=4,
+      num_groups=32,
+      param_dtype=float32,
+      reduction_axes=None,
+      scale=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      scale_init=<function ones at 0x7fb32b86e520>,
+      use_bias=True,
+      use_fast_variance=True,
+      use_scale=True
+    ),
+    out_channels=128,
+    time_mlp=Linear( # Param: 16,512 (66.0 KB)
+      bias=Param( # 128 (512 B)
+        value=Array(shape=(128,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dot_general=<function dot_general at 0x7fb32c21a3e0>,
+      dtype=float32,
+      in_features=256,
+      kernel=Param( # 16,384 (65.5 KB)
+        value=Array(shape=(128, 128), dtype=dtype('float32'))
+      ),
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      out_features=256,
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      use_bias=True
+    )
+  )]
+)],
+    blocks_up=[TimestepEmbedSequential( # Param: 135,872 (543.5 KB), RngState: 2 (12 B), Total: 135,874 (543.5 KB)
+  layers=[ResnetBlock( # Param: 135,872 (543.5 KB), RngState: 2 (12 B), Total: 135,874 (543.5 KB)
+    activation=<PjitFunction of <function silu at 0x7fb32b8440e0>>,
+    conv1=Conv( # Param: 73,792 (295.2 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=256,
+      input_dilation=1,
+      kernel=Param( # 73,728 (294.9 KB)
+        value=Array(shape=(3, 3, 128, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(3, 3, 256, 128),
+      kernel_size=(3, 3),
+      mask=None,
+      out_features=128,
+      padding=(1, 1),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    conv2=Conv( # Param: 36,928 (147.7 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 36,864 (147.5 KB)
+        value=Array(shape=(3, 3, 64, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(3, 3, 128, 128),
+      kernel_size=(3, 3),
+      mask=None,
+      out_features=128,
+      padding=(1, 1),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    dropout=Dropout( # RngState: 2 (12 B)
+      broadcast_dims=(),
+      deterministic=False,
+      rate=0.1,
+      rng_collection='dropout',
+      rngs=RngStream( # RngState: 2 (12 B)
+        count=RngCount( # 1 (4 B)
+          value=Array(0, dtype=uint32),
+          tag='default'
+        ),
+        key=RngKey( # 1 (8 B)
+          value=Array((), dtype=key<fry>) overlaying:
+          [2853902436 2217684095],
+          tag='default'
+        ),
+        tag='default'
+      )
+    ),
+    embedding_dim=256,
+    in_channels=256,
+    nin_shortcut=Conv( # Param: 8,256 (33.0 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=256,
+      input_dilation=1,
+      kernel=Param( # 8,192 (32.8 KB)
+        value=Array(shape=(1, 1, 128, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(1, 1, 256, 128),
+      kernel_size=(1, 1),
+      mask=None,
+      out_features=128,
+      padding=(0, 0),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    norm1=GroupNorm( # Param: 256 (1.0 KB)
+      axis_index_groups=None,
+      axis_name=None,
+      bias=Param( # 128 (512 B)
+        value=Array(shape=(128,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dtype=float32,
+      epsilon=1e-06,
+      feature_axis=-1,
+      group_size=8,
+      num_groups=32,
+      param_dtype=float32,
+      reduction_axes=None,
+      scale=Param( # 128 (512 B)
+        value=Array(shape=(128,), dtype=dtype('float32'))
+      ),
+      scale_init=<function ones at 0x7fb32b86e520>,
+      use_bias=True,
+      use_fast_variance=True,
+      use_scale=True
+    ),
+    norm2=GroupNorm( # Param: 128 (512 B)
+      axis_index_groups=None,
+      axis_name=None,
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dtype=float32,
+      epsilon=1e-06,
+      feature_axis=-1,
+      group_size=4,
+      num_groups=32,
+      param_dtype=float32,
+      reduction_axes=None,
+      scale=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      scale_init=<function ones at 0x7fb32b86e520>,
+      use_bias=True,
+      use_fast_variance=True,
+      use_scale=True
+    ),
+    out_channels=128,
+    time_mlp=Linear( # Param: 16,512 (66.0 KB)
+      bias=Param( # 128 (512 B)
+        value=Array(shape=(128,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dot_general=<function dot_general at 0x7fb32c21a3e0>,
+      dtype=float32,
+      in_features=256,
+      kernel=Param( # 16,384 (65.5 KB)
+        value=Array(shape=(128, 128), dtype=dtype('float32'))
+      ),
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      out_features=256,
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      use_bias=True
+    )
+  )]
+), TimestepEmbedSequential( # Param: 320,512 (1.3 MB), RngState: 2 (12 B), Total: 320,514 (1.3 MB)
+  layers=[ResnetBlock( # Param: 135,872 (543.5 KB), RngState: 2 (12 B), Total: 135,874 (543.5 KB)
+    activation=<PjitFunction of <function silu at 0x7fb32b8440e0>>,
+    conv1=Conv( # Param: 73,792 (295.2 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=256,
+      input_dilation=1,
+      kernel=Param( # 73,728 (294.9 KB)
+        value=Array(shape=(3, 3, 128, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(3, 3, 256, 128),
+      kernel_size=(3, 3),
+      mask=None,
+      out_features=128,
+      padding=(1, 1),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    conv2=Conv( # Param: 36,928 (147.7 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 36,864 (147.5 KB)
+        value=Array(shape=(3, 3, 64, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(3, 3, 128, 128),
+      kernel_size=(3, 3),
+      mask=None,
+      out_features=128,
+      padding=(1, 1),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    dropout=Dropout( # RngState: 2 (12 B)
+      broadcast_dims=(),
+      deterministic=False,
+      rate=0.1,
+      rng_collection='dropout',
+      rngs=RngStream( # RngState: 2 (12 B)
+        count=RngCount( # 1 (4 B)
+          value=Array(0, dtype=uint32),
+          tag='default'
+        ),
+        key=RngKey( # 1 (8 B)
+          value=Array((), dtype=key<fry>) overlaying:
+          [2785098898  841100811],
+          tag='default'
+        ),
+        tag='default'
+      )
+    ),
+    embedding_dim=256,
+    in_channels=256,
+    nin_shortcut=Conv( # Param: 8,256 (33.0 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=256,
+      input_dilation=1,
+      kernel=Param( # 8,192 (32.8 KB)
+        value=Array(shape=(1, 1, 128, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(1, 1, 256, 128),
+      kernel_size=(1, 1),
+      mask=None,
+      out_features=128,
+      padding=(0, 0),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    norm1=GroupNorm( # Param: 256 (1.0 KB)
+      axis_index_groups=None,
+      axis_name=None,
+      bias=Param( # 128 (512 B)
+        value=Array(shape=(128,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dtype=float32,
+      epsilon=1e-06,
+      feature_axis=-1,
+      group_size=8,
+      num_groups=32,
+      param_dtype=float32,
+      reduction_axes=None,
+      scale=Param( # 128 (512 B)
+        value=Array(shape=(128,), dtype=dtype('float32'))
+      ),
+      scale_init=<function ones at 0x7fb32b86e520>,
+      use_bias=True,
+      use_fast_variance=True,
+      use_scale=True
+    ),
+    norm2=GroupNorm( # Param: 128 (512 B)
+      axis_index_groups=None,
+      axis_name=None,
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dtype=float32,
+      epsilon=1e-06,
+      feature_axis=-1,
+      group_size=4,
+      num_groups=32,
+      param_dtype=float32,
+      reduction_axes=None,
+      scale=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      scale_init=<function ones at 0x7fb32b86e520>,
+      use_bias=True,
+      use_fast_variance=True,
+      use_scale=True
+    ),
+    out_channels=128,
+    time_mlp=Linear( # Param: 16,512 (66.0 KB)
+      bias=Param( # 128 (512 B)
+        value=Array(shape=(128,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dot_general=<function dot_general at 0x7fb32c21a3e0>,
+      dtype=float32,
+      in_features=256,
+      kernel=Param( # 16,384 (65.5 KB)
+        value=Array(shape=(128, 128), dtype=dtype('float32'))
+      ),
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      out_features=256,
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      use_bias=True
+    )
+  ), Upsample( # Param: 184,640 (738.6 KB)
+    conv_pixel_shuffle=Conv( # Param: 147,712 (590.8 KB)
+      bias=Param( # 256 (1.0 KB)
+        value=Array(shape=(256,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 147,456 (589.8 KB)
+        value=Array(shape=(3, 3, 64, 256), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(3, 3, 128, 512),
+      kernel_size=(3, 3),
+      mask=None,
+      out_features=512,
+      padding=(1, 1),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    conv_resize=Conv( # Param: 36,928 (147.7 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 36,864 (147.5 KB)
+        value=Array(shape=(3, 3, 64, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(3, 3, 128, 128),
+      kernel_size=(3, 3),
+      mask=None,
+      out_features=128,
+      padding=(1, 1),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    method='pixel_shuffle',
+    pixel_shuffle=PixelShuffle(
+      scale=2
+    ),
+    scale_factor=2
+  )]
+), TimestepEmbedSequential( # Param: 152,640 (610.6 KB), RngState: 2 (12 B), Total: 152,642 (610.6 KB)
+  layers=[ResnetBlock( # Param: 135,872 (543.5 KB), RngState: 2 (12 B), Total: 135,874 (543.5 KB)
+    activation=<PjitFunction of <function silu at 0x7fb32b8440e0>>,
+    conv1=Conv( # Param: 73,792 (295.2 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=256,
+      input_dilation=1,
+      kernel=Param( # 73,728 (294.9 KB)
+        value=Array(shape=(3, 3, 128, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(3, 3, 256, 128),
+      kernel_size=(3, 3),
+      mask=None,
+      out_features=128,
+      padding=(1, 1),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    conv2=Conv( # Param: 36,928 (147.7 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 36,864 (147.5 KB)
+        value=Array(shape=(3, 3, 64, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(3, 3, 128, 128),
+      kernel_size=(3, 3),
+      mask=None,
+      out_features=128,
+      padding=(1, 1),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    dropout=Dropout( # RngState: 2 (12 B)
+      broadcast_dims=(),
+      deterministic=False,
+      rate=0.1,
+      rng_collection='dropout',
+      rngs=RngStream( # RngState: 2 (12 B)
+        count=RngCount( # 1 (4 B)
+          value=Array(0, dtype=uint32),
+          tag='default'
+        ),
+        key=RngKey( # 1 (8 B)
+          value=Array((), dtype=key<fry>) overlaying:
+          [  48802331 1548237274],
+          tag='default'
+        ),
+        tag='default'
+      )
+    ),
+    embedding_dim=256,
+    in_channels=256,
+    nin_shortcut=Conv( # Param: 8,256 (33.0 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=256,
+      input_dilation=1,
+      kernel=Param( # 8,192 (32.8 KB)
+        value=Array(shape=(1, 1, 128, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(1, 1, 256, 128),
+      kernel_size=(1, 1),
+      mask=None,
+      out_features=128,
+      padding=(0, 0),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    norm1=GroupNorm( # Param: 256 (1.0 KB)
+      axis_index_groups=None,
+      axis_name=None,
+      bias=Param( # 128 (512 B)
+        value=Array(shape=(128,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dtype=float32,
+      epsilon=1e-06,
+      feature_axis=-1,
+      group_size=8,
+      num_groups=32,
+      param_dtype=float32,
+      reduction_axes=None,
+      scale=Param( # 128 (512 B)
+        value=Array(shape=(128,), dtype=dtype('float32'))
+      ),
+      scale_init=<function ones at 0x7fb32b86e520>,
+      use_bias=True,
+      use_fast_variance=True,
+      use_scale=True
+    ),
+    norm2=GroupNorm( # Param: 128 (512 B)
+      axis_index_groups=None,
+      axis_name=None,
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dtype=float32,
+      epsilon=1e-06,
+      feature_axis=-1,
+      group_size=4,
+      num_groups=32,
+      param_dtype=float32,
+      reduction_axes=None,
+      scale=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      scale_init=<function ones at 0x7fb32b86e520>,
+      use_bias=True,
+      use_fast_variance=True,
+      use_scale=True
+    ),
+    out_channels=128,
+    time_mlp=Linear( # Param: 16,512 (66.0 KB)
+      bias=Param( # 128 (512 B)
+        value=Array(shape=(128,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dot_general=<function dot_general at 0x7fb32c21a3e0>,
+      dtype=float32,
+      in_features=256,
+      kernel=Param( # 16,384 (65.5 KB)
+        value=Array(shape=(128, 128), dtype=dtype('float32'))
+      ),
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      out_features=256,
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      use_bias=True
+    )
+  ), AttnBlock( # Param: 16,768 (67.1 KB)
+    dtype=float32,
+    head_dim=32,
+    k=Conv( # Param: 4,160 (16.6 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 4,096 (16.4 KB)
+        value=Array(shape=(1, 1, 64, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(1, 1, 128, 128),
+      kernel_size=(1, 1),
+      mask=None,
+      out_features=128,
+      padding='SAME',
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=1,
+      use_bias=True
+    ),
+    norm=GroupNorm( # Param: 128 (512 B)
+      axis_index_groups=None,
+      axis_name=None,
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dtype=float32,
+      epsilon=1e-06,
+      feature_axis=-1,
+      group_size=4,
+      num_groups=32,
+      param_dtype=float32,
+      reduction_axes=None,
+      scale=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      scale_init=<function ones at 0x7fb32b86e520>,
+      use_bias=True,
+      use_fast_variance=True,
+      use_scale=True
+    ),
+    num_heads=4,
+    proj_out=Conv( # Param: 4,160 (16.6 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 4,096 (16.4 KB)
+        value=Array(shape=(1, 1, 64, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(1, 1, 128, 128),
+      kernel_size=(1, 1),
+      mask=None,
+      out_features=128,
+      padding='SAME',
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=1,
+      use_bias=True
+    ),
+    q=Conv( # Param: 4,160 (16.6 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 4,096 (16.4 KB)
+        value=Array(shape=(1, 1, 64, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(1, 1, 128, 128),
+      kernel_size=(1, 1),
+      mask=None,
+      out_features=128,
+      padding='SAME',
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=1,
+      use_bias=True
+    ),
+    v=Conv( # Param: 4,160 (16.6 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 4,096 (16.4 KB)
+        value=Array(shape=(1, 1, 64, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(1, 1, 128, 128),
+      kernel_size=(1, 1),
+      mask=None,
+      out_features=128,
+      padding='SAME',
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=1,
+      use_bias=True
+    )
+  )]
+), TimestepEmbedSequential( # Param: 316,736 (1.3 MB), RngState: 2 (12 B), Total: 316,738 (1.3 MB)
+  layers=[ResnetBlock( # Param: 115,328 (461.3 KB), RngState: 2 (12 B), Total: 115,330 (461.3 KB)
+    activation=<PjitFunction of <function silu at 0x7fb32b8440e0>>,
+    conv1=Conv( # Param: 55,360 (221.4 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=192,
+      input_dilation=1,
+      kernel=Param( # 55,296 (221.2 KB)
+        value=Array(shape=(3, 3, 96, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(3, 3, 192, 128),
+      kernel_size=(3, 3),
+      mask=None,
+      out_features=128,
+      padding=(1, 1),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    conv2=Conv( # Param: 36,928 (147.7 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 36,864 (147.5 KB)
+        value=Array(shape=(3, 3, 64, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(3, 3, 128, 128),
+      kernel_size=(3, 3),
+      mask=None,
+      out_features=128,
+      padding=(1, 1),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    dropout=Dropout( # RngState: 2 (12 B)
+      broadcast_dims=(),
+      deterministic=False,
+      rate=0.1,
+      rng_collection='dropout',
+      rngs=RngStream( # RngState: 2 (12 B)
+        count=RngCount( # 1 (4 B)
+          value=Array(0, dtype=uint32),
+          tag='default'
+        ),
+        key=RngKey( # 1 (8 B)
+          value=Array((), dtype=key<fry>) overlaying:
+          [1596966061 1315822572],
+          tag='default'
+        ),
+        tag='default'
+      )
+    ),
+    embedding_dim=256,
+    in_channels=192,
+    nin_shortcut=Conv( # Param: 6,208 (24.8 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=192,
+      input_dilation=1,
+      kernel=Param( # 6,144 (24.6 KB)
+        value=Array(shape=(1, 1, 96, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(1, 1, 192, 128),
+      kernel_size=(1, 1),
+      mask=None,
+      out_features=128,
+      padding=(0, 0),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    norm1=GroupNorm( # Param: 192 (768 B)
+      axis_index_groups=None,
+      axis_name=None,
+      bias=Param( # 96 (384 B)
+        value=Array(shape=(96,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dtype=float32,
+      epsilon=1e-06,
+      feature_axis=-1,
+      group_size=6,
+      num_groups=32,
+      param_dtype=float32,
+      reduction_axes=None,
+      scale=Param( # 96 (384 B)
+        value=Array(shape=(96,), dtype=dtype('float32'))
+      ),
+      scale_init=<function ones at 0x7fb32b86e520>,
+      use_bias=True,
+      use_fast_variance=True,
+      use_scale=True
+    ),
+    norm2=GroupNorm( # Param: 128 (512 B)
+      axis_index_groups=None,
+      axis_name=None,
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dtype=float32,
+      epsilon=1e-06,
+      feature_axis=-1,
+      group_size=4,
+      num_groups=32,
+      param_dtype=float32,
+      reduction_axes=None,
+      scale=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      scale_init=<function ones at 0x7fb32b86e520>,
+      use_bias=True,
+      use_fast_variance=True,
+      use_scale=True
+    ),
+    out_channels=128,
+    time_mlp=Linear( # Param: 16,512 (66.0 KB)
+      bias=Param( # 128 (512 B)
+        value=Array(shape=(128,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dot_general=<function dot_general at 0x7fb32c21a3e0>,
+      dtype=float32,
+      in_features=256,
+      kernel=Param( # 16,384 (65.5 KB)
+        value=Array(shape=(128, 128), dtype=dtype('float32'))
+      ),
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      out_features=256,
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      use_bias=True
+    )
+  ), AttnBlock( # Param: 16,768 (67.1 KB)
+    dtype=float32,
+    head_dim=32,
+    k=Conv( # Param: 4,160 (16.6 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 4,096 (16.4 KB)
+        value=Array(shape=(1, 1, 64, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(1, 1, 128, 128),
+      kernel_size=(1, 1),
+      mask=None,
+      out_features=128,
+      padding='SAME',
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=1,
+      use_bias=True
+    ),
+    norm=GroupNorm( # Param: 128 (512 B)
+      axis_index_groups=None,
+      axis_name=None,
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dtype=float32,
+      epsilon=1e-06,
+      feature_axis=-1,
+      group_size=4,
+      num_groups=32,
+      param_dtype=float32,
+      reduction_axes=None,
+      scale=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      scale_init=<function ones at 0x7fb32b86e520>,
+      use_bias=True,
+      use_fast_variance=True,
+      use_scale=True
+    ),
+    num_heads=4,
+    proj_out=Conv( # Param: 4,160 (16.6 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 4,096 (16.4 KB)
+        value=Array(shape=(1, 1, 64, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(1, 1, 128, 128),
+      kernel_size=(1, 1),
+      mask=None,
+      out_features=128,
+      padding='SAME',
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=1,
+      use_bias=True
+    ),
+    q=Conv( # Param: 4,160 (16.6 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 4,096 (16.4 KB)
+        value=Array(shape=(1, 1, 64, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(1, 1, 128, 128),
+      kernel_size=(1, 1),
+      mask=None,
+      out_features=128,
+      padding='SAME',
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=1,
+      use_bias=True
+    ),
+    v=Conv( # Param: 4,160 (16.6 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 4,096 (16.4 KB)
+        value=Array(shape=(1, 1, 64, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(1, 1, 128, 128),
+      kernel_size=(1, 1),
+      mask=None,
+      out_features=128,
+      padding='SAME',
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=1,
+      use_bias=True
+    )
+  ), Upsample( # Param: 184,640 (738.6 KB)
+    conv_pixel_shuffle=Conv( # Param: 147,712 (590.8 KB)
+      bias=Param( # 256 (1.0 KB)
+        value=Array(shape=(256,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 147,456 (589.8 KB)
+        value=Array(shape=(3, 3, 64, 256), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(3, 3, 128, 512),
+      kernel_size=(3, 3),
+      mask=None,
+      out_features=512,
+      padding=(1, 1),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    conv_resize=Conv( # Param: 36,928 (147.7 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 36,864 (147.5 KB)
+        value=Array(shape=(3, 3, 64, 64), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(3, 3, 128, 128),
+      kernel_size=(3, 3),
+      mask=None,
+      out_features=128,
+      padding=(1, 1),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    method='pixel_shuffle',
+    pixel_shuffle=PixelShuffle(
+      scale=2
+    ),
+    scale_factor=2
+  )]
+), TimestepEmbedSequential( # Param: 48,544 (194.2 KB), RngState: 2 (12 B), Total: 48,546 (194.2 KB)
+  layers=[ResnetBlock( # Param: 48,544 (194.2 KB), RngState: 2 (12 B), Total: 48,546 (194.2 KB)
+    activation=<PjitFunction of <function silu at 0x7fb32b8440e0>>,
+    conv1=Conv( # Param: 27,680 (110.7 KB)
+      bias=Param( # 32 (128 B)
+        value=Array(shape=(32,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=192,
+      input_dilation=1,
+      kernel=Param( # 27,648 (110.6 KB)
+        value=Array(shape=(3, 3, 96, 32), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(3, 3, 192, 64),
+      kernel_size=(3, 3),
+      mask=None,
+      out_features=64,
+      padding=(1, 1),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    conv2=Conv( # Param: 9,248 (37.0 KB)
+      bias=Param( # 32 (128 B)
+        value=Array(shape=(32,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=64,
+      input_dilation=1,
+      kernel=Param( # 9,216 (36.9 KB)
+        value=Array(shape=(3, 3, 32, 32), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(3, 3, 64, 64),
+      kernel_size=(3, 3),
+      mask=None,
+      out_features=64,
+      padding=(1, 1),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    dropout=Dropout( # RngState: 2 (12 B)
+      broadcast_dims=(),
+      deterministic=False,
+      rate=0.1,
+      rng_collection='dropout',
+      rngs=RngStream( # RngState: 2 (12 B)
+        count=RngCount( # 1 (4 B)
+          value=Array(0, dtype=uint32),
+          tag='default'
+        ),
+        key=RngKey( # 1 (8 B)
+          value=Array((), dtype=key<fry>) overlaying:
+          [2550820645 2818876438],
+          tag='default'
+        ),
+        tag='default'
+      )
+    ),
+    embedding_dim=256,
+    in_channels=192,
+    nin_shortcut=Conv( # Param: 3,104 (12.4 KB)
+      bias=Param( # 32 (128 B)
+        value=Array(shape=(32,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=192,
+      input_dilation=1,
+      kernel=Param( # 3,072 (12.3 KB)
+        value=Array(shape=(1, 1, 96, 32), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(1, 1, 192, 64),
+      kernel_size=(1, 1),
+      mask=None,
+      out_features=64,
+      padding=(0, 0),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    norm1=GroupNorm( # Param: 192 (768 B)
+      axis_index_groups=None,
+      axis_name=None,
+      bias=Param( # 96 (384 B)
+        value=Array(shape=(96,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dtype=float32,
+      epsilon=1e-06,
+      feature_axis=-1,
+      group_size=6,
+      num_groups=32,
+      param_dtype=float32,
+      reduction_axes=None,
+      scale=Param( # 96 (384 B)
+        value=Array(shape=(96,), dtype=dtype('float32'))
+      ),
+      scale_init=<function ones at 0x7fb32b86e520>,
+      use_bias=True,
+      use_fast_variance=True,
+      use_scale=True
+    ),
+    norm2=GroupNorm( # Param: 64 (256 B)
+      axis_index_groups=None,
+      axis_name=None,
+      bias=Param( # 32 (128 B)
+        value=Array(shape=(32,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dtype=float32,
+      epsilon=1e-06,
+      feature_axis=-1,
+      group_size=2,
+      num_groups=32,
+      param_dtype=float32,
+      reduction_axes=None,
+      scale=Param( # 32 (128 B)
+        value=Array(shape=(32,), dtype=dtype('float32'))
+      ),
+      scale_init=<function ones at 0x7fb32b86e520>,
+      use_bias=True,
+      use_fast_variance=True,
+      use_scale=True
+    ),
+    out_channels=64,
+    time_mlp=Linear( # Param: 8,256 (33.0 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dot_general=<function dot_general at 0x7fb32c21a3e0>,
+      dtype=float32,
+      in_features=256,
+      kernel=Param( # 8,192 (32.8 KB)
+        value=Array(shape=(128, 64), dtype=dtype('float32'))
+      ),
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      out_features=128,
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      use_bias=True
+    )
+  )]
+), TimestepEmbedSequential( # Param: 38,240 (153.0 KB), RngState: 2 (12 B), Total: 38,242 (153.0 KB)
+  layers=[ResnetBlock( # Param: 38,240 (153.0 KB), RngState: 2 (12 B), Total: 38,242 (153.0 KB)
+    activation=<PjitFunction of <function silu at 0x7fb32b8440e0>>,
+    conv1=Conv( # Param: 18,464 (73.9 KB)
+      bias=Param( # 32 (128 B)
+        value=Array(shape=(32,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 18,432 (73.7 KB)
+        value=Array(shape=(3, 3, 64, 32), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(3, 3, 128, 64),
+      kernel_size=(3, 3),
+      mask=None,
+      out_features=64,
+      padding=(1, 1),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    conv2=Conv( # Param: 9,248 (37.0 KB)
+      bias=Param( # 32 (128 B)
+        value=Array(shape=(32,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=64,
+      input_dilation=1,
+      kernel=Param( # 9,216 (36.9 KB)
+        value=Array(shape=(3, 3, 32, 32), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(3, 3, 64, 64),
+      kernel_size=(3, 3),
+      mask=None,
+      out_features=64,
+      padding=(1, 1),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    dropout=Dropout( # RngState: 2 (12 B)
+      broadcast_dims=(),
+      deterministic=False,
+      rate=0.1,
+      rng_collection='dropout',
+      rngs=RngStream( # RngState: 2 (12 B)
+        count=RngCount( # 1 (4 B)
+          value=Array(0, dtype=uint32),
+          tag='default'
+        ),
+        key=RngKey( # 1 (8 B)
+          value=Array((), dtype=key<fry>) overlaying:
+          [1975238715 3717004500],
+          tag='default'
+        ),
+        tag='default'
+      )
+    ),
+    embedding_dim=256,
+    in_channels=128,
+    nin_shortcut=Conv( # Param: 2,080 (8.3 KB)
+      bias=Param( # 32 (128 B)
+        value=Array(shape=(32,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      conv_general_dilated=<function conv_general_dilated at 0x7fb32c260ea0>,
+      dtype=float32,
+      feature_group_count=1,
+      in_features=128,
+      input_dilation=1,
+      kernel=Param( # 2,048 (8.2 KB)
+        value=Array(shape=(1, 1, 64, 32), dtype=dtype('float32'))
+      ),
+      kernel_dilation=1,
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      kernel_shape=(1, 1, 128, 64),
+      kernel_size=(1, 1),
+      mask=None,
+      out_features=64,
+      padding=(0, 0),
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      strides=(1, 1),
+      use_bias=True
+    ),
+    norm1=GroupNorm( # Param: 128 (512 B)
+      axis_index_groups=None,
+      axis_name=None,
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dtype=float32,
+      epsilon=1e-06,
+      feature_axis=-1,
+      group_size=4,
+      num_groups=32,
+      param_dtype=float32,
+      reduction_axes=None,
+      scale=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      scale_init=<function ones at 0x7fb32b86e520>,
+      use_bias=True,
+      use_fast_variance=True,
+      use_scale=True
+    ),
+    norm2=GroupNorm( # Param: 64 (256 B)
+      axis_index_groups=None,
+      axis_name=None,
+      bias=Param( # 32 (128 B)
+        value=Array(shape=(32,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dtype=float32,
+      epsilon=1e-06,
+      feature_axis=-1,
+      group_size=2,
+      num_groups=32,
+      param_dtype=float32,
+      reduction_axes=None,
+      scale=Param( # 32 (128 B)
+        value=Array(shape=(32,), dtype=dtype('float32'))
+      ),
+      scale_init=<function ones at 0x7fb32b86e520>,
+      use_bias=True,
+      use_fast_variance=True,
+      use_scale=True
+    ),
+    out_channels=64,
+    time_mlp=Linear( # Param: 8,256 (33.0 KB)
+      bias=Param( # 64 (256 B)
+        value=Array(shape=(64,), dtype=dtype('float32'))
+      ),
+      bias_init=<function zeros at 0x7fb32b98c2c0>,
+      dot_general=<function dot_general at 0x7fb32c21a3e0>,
+      dtype=float32,
+      in_features=256,
+      kernel=Param( # 8,192 (32.8 KB)
+        value=Array(shape=(128, 64), dtype=dtype('float32'))
+      ),
+      kernel_init=<function variance_scaling.<locals>.init at 0x7fb3299d7060>,
+      out_features=128,
+      param_dtype=float32,
+      precision=None,
+      promote_dtype=<function promote_dtype at 0x7fb3299d7380>,
+      use_bias=True
+    )
+  )]
+)],
 )
-
-# Metadata
-STEP = 45000
-PRECISION = "float32"
-MODEL_TYPE = "CondUNet2D"