KRadim's picture
Upload Conditional MediumUNet model (Orbax checkpoint + native MLIR + compiled StableHLO artifact)
19ff442 verified
Raw
History Blame Contribute Delete
24.5 kB
{"tree_metadata": {"('class_emb', 'embedding', 'value')": {"key_metadata": [{"key": "class_emb", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [201, 64]}}, "('cond_proj1', 'bias', 'value')": {"key_metadata": [{"key": "cond_proj1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('cond_proj1', 'kernel', 'value')": {"key_metadata": [{"key": "cond_proj1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('cond_proj2', 'bias', 'value')": {"key_metadata": [{"key": "cond_proj2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('cond_proj2', 'kernel', 'value')": {"key_metadata": [{"key": "cond_proj2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('down1_block1', 'cond_proj', 'bias', 'value')": {"key_metadata": [{"key": "down1_block1", "key_type": 2}, {"key": "cond_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('down1_block1', 'cond_proj', 'kernel', 'value')": {"key_metadata": [{"key": "down1_block1", "key_type": 2}, {"key": "cond_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 32]}}, "('down1_block1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "down1_block1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('down1_block1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "down1_block1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 3, 32]}}, "('down1_block1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "down1_block1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('down1_block1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "down1_block1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 32, 32]}}, "('down1_block1', 'gn1', 'bias', 'value')": {"key_metadata": [{"key": "down1_block1", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('down1_block1', 'gn1', 'scale', 'value')": {"key_metadata": [{"key": "down1_block1", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('down1_block1', 'gn2', 'bias', 'value')": {"key_metadata": [{"key": "down1_block1", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('down1_block1', 'gn2', 'scale', 'value')": {"key_metadata": [{"key": "down1_block1", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('down1_block1', 'skip_conv', 'bias', 'value')": {"key_metadata": [{"key": "down1_block1", "key_type": 2}, {"key": "skip_conv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('down1_block1', 'skip_conv', 'kernel', 'value')": {"key_metadata": [{"key": "down1_block1", "key_type": 2}, {"key": "skip_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 3, 32]}}, "('down1_block2', 'cond_proj', 'bias', 'value')": {"key_metadata": [{"key": "down1_block2", "key_type": 2}, {"key": "cond_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('down1_block2', 'cond_proj', 'kernel', 'value')": {"key_metadata": [{"key": "down1_block2", "key_type": 2}, {"key": "cond_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 32]}}, "('down1_block2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "down1_block2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('down1_block2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "down1_block2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 32, 32]}}, "('down1_block2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "down1_block2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('down1_block2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "down1_block2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 32, 32]}}, "('down1_block2', 'gn1', 'bias', 'value')": {"key_metadata": [{"key": "down1_block2", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('down1_block2', 'gn1', 'scale', 'value')": {"key_metadata": [{"key": "down1_block2", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('down1_block2', 'gn2', 'bias', 'value')": {"key_metadata": [{"key": "down1_block2", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('down1_block2', 'gn2', 'scale', 'value')": {"key_metadata": [{"key": "down1_block2", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('down2_block1', 'cond_proj', 'bias', 'value')": {"key_metadata": [{"key": "down2_block1", "key_type": 2}, {"key": "cond_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('down2_block1', 'cond_proj', 'kernel', 'value')": {"key_metadata": [{"key": "down2_block1", "key_type": 2}, {"key": "cond_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 64]}}, "('down2_block1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "down2_block1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('down2_block1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "down2_block1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('down2_block1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "down2_block1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('down2_block1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "down2_block1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('down2_block1', 'gn1', 'bias', 'value')": {"key_metadata": [{"key": "down2_block1", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('down2_block1', 'gn1', 'scale', 'value')": {"key_metadata": [{"key": "down2_block1", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('down2_block1', 'gn2', 'bias', 'value')": {"key_metadata": [{"key": "down2_block1", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('down2_block1', 'gn2', 'scale', 'value')": {"key_metadata": [{"key": "down2_block1", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('down2_stride', 'bias', 'value')": {"key_metadata": [{"key": "down2_stride", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('down2_stride', 'kernel', 'value')": {"key_metadata": [{"key": "down2_stride", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 32, 64]}}, "('down3_block1', 'cond_proj', 'bias', 'value')": {"key_metadata": [{"key": "down3_block1", "key_type": 2}, {"key": "cond_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('down3_block1', 'cond_proj', 'kernel', 'value')": {"key_metadata": [{"key": "down3_block1", "key_type": 2}, {"key": "cond_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('down3_block1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "down3_block1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('down3_block1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "down3_block1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('down3_block1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "down3_block1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('down3_block1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "down3_block1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('down3_block1', 'gn1', 'bias', 'value')": {"key_metadata": [{"key": "down3_block1", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('down3_block1', 'gn1', 'scale', 'value')": {"key_metadata": [{"key": "down3_block1", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('down3_block1', 'gn2', 'bias', 'value')": {"key_metadata": [{"key": "down3_block1", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('down3_block1', 'gn2', 'scale', 'value')": {"key_metadata": [{"key": "down3_block1", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('down3_stride', 'bias', 'value')": {"key_metadata": [{"key": "down3_stride", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('down3_stride', 'kernel', 'value')": {"key_metadata": [{"key": "down3_stride", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 128]}}, "('out_conv', 'bias', 'value')": {"key_metadata": [{"key": "out_conv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('out_conv', 'kernel', 'value')": {"key_metadata": [{"key": "out_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 32, 3]}}, "('out_gn', 'bias', 'value')": {"key_metadata": [{"key": "out_gn", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('out_gn', 'scale', 'value')": {"key_metadata": [{"key": "out_gn", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('up1_block1', 'cond_proj', 'bias', 'value')": {"key_metadata": [{"key": "up1_block1", "key_type": 2}, {"key": "cond_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('up1_block1', 'cond_proj', 'kernel', 'value')": {"key_metadata": [{"key": "up1_block1", "key_type": 2}, {"key": "cond_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 64]}}, "('up1_block1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "up1_block1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('up1_block1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "up1_block1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 64]}}, "('up1_block1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "up1_block1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('up1_block1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "up1_block1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('up1_block1', 'gn1', 'bias', 'value')": {"key_metadata": [{"key": "up1_block1", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('up1_block1', 'gn1', 'scale', 'value')": {"key_metadata": [{"key": "up1_block1", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('up1_block1', 'gn2', 'bias', 'value')": {"key_metadata": [{"key": "up1_block1", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('up1_block1', 'gn2', 'scale', 'value')": {"key_metadata": [{"key": "up1_block1", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('up1_block1', 'skip_conv', 'bias', 'value')": {"key_metadata": [{"key": "up1_block1", "key_type": 2}, {"key": "skip_conv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('up1_block1', 'skip_conv', 'kernel', 'value')": {"key_metadata": [{"key": "up1_block1", "key_type": 2}, {"key": "skip_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 64]}}, "('up1_transpose', 'bias', 'value')": {"key_metadata": [{"key": "up1_transpose", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('up1_transpose', 'kernel', 'value')": {"key_metadata": [{"key": "up1_transpose", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 64]}}, "('up2_block1', 'cond_proj', 'bias', 'value')": {"key_metadata": [{"key": "up2_block1", "key_type": 2}, {"key": "cond_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('up2_block1', 'cond_proj', 'kernel', 'value')": {"key_metadata": [{"key": "up2_block1", "key_type": 2}, {"key": "cond_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 32]}}, "('up2_block1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "up2_block1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('up2_block1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "up2_block1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 32]}}, "('up2_block1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "up2_block1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('up2_block1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "up2_block1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 32, 32]}}, "('up2_block1', 'gn1', 'bias', 'value')": {"key_metadata": [{"key": "up2_block1", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('up2_block1', 'gn1', 'scale', 'value')": {"key_metadata": [{"key": "up2_block1", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('up2_block1', 'gn2', 'bias', 'value')": {"key_metadata": [{"key": "up2_block1", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('up2_block1', 'gn2', 'scale', 'value')": {"key_metadata": [{"key": "up2_block1", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('up2_block1', 'skip_conv', 'bias', 'value')": {"key_metadata": [{"key": "up2_block1", "key_type": 2}, {"key": "skip_conv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('up2_block1', 'skip_conv', 'kernel', 'value')": {"key_metadata": [{"key": "up2_block1", "key_type": 2}, {"key": "skip_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 64, 32]}}, "('up2_transpose', 'bias', 'value')": {"key_metadata": [{"key": "up2_transpose", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('up2_transpose', 'kernel', 'value')": {"key_metadata": [{"key": "up2_transpose", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 32]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}