Upload model at step 45000

config.py

@@ -1,2042 +1,26 @@
 """Model configuration for jcopo/mnist
 
-This file contains the model architecture definition.
 Training step: 45000
 Precision: float32
 """
 
-from triax.models.nn.condUNet import CondUNet2D
+import jax
 import jax.numpy as jnp
 from flax import nnx
+from triax.models import CondUNet2D
 
-# Model architecture
+# Model definition
 model = CondUNet2D(
-    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
+        in_channels=IN_CHANNELS,
+        ch=variant["base_channels"],
+        ch_mult=variant["channel_multipliers"],
+        num_res_blocks=variant["num_res_blocks"],
+        attention_resolutions=variant["attention_resolutions"],
+        num_heads=variant["num_heads"],
+        dropout=DROPOUT,
+        dropout_rate=DROPOUT_RATE,
+        activation=nnx.swish,
+        param_dtype=jnp.float32,
+        dtype=compute_dtype,
+        rngs=nnx.Rngs(MODEL_SEED),
     )
-  )]
-), 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
-    )
-  )]
-)],
-)