{"tree_metadata": {"('step',)": {"key_metadata": [{"key": "step", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'Conv_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'DownSample_0', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "DownSample_0", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'DownSample_0', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "DownSample_0", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'DownSample_1', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "DownSample_1", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'DownSample_1', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "DownSample_1", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'DownSample_2', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "DownSample_2", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'DownSample_2', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "DownSample_2", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_0', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_0', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_1', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_1', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_2', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_2', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_3', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_3', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_4', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_4', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_5', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_5', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_0', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_0', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_0', 'Conv_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_0', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_0', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_0', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_0', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_0', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_1', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_1', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_1', 'Conv_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_1', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_1', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_1', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_1', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_1', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_1', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_1', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_10', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_10', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_10', 'Conv_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_10', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_10', 'Conv_2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_10', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_10', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_10', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_10', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_10', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_10', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_10', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_11', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_11', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_11', 'Conv_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_11', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_11', 'Conv_2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_11', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_11', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_11', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_11', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_11', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_11', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_11', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_12', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_12', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_12', 'Conv_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_12', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_12', 'Conv_2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_12', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_12', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_12', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_12', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_12', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_12', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_12', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_13', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_13', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_13', 'Conv_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_13', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_13', 'Conv_2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_13', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_13', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_13', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_13', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_13', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_13', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_13', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_14', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_14', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_14', 'Conv_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_14', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_14', 'Conv_2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_14', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_14', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_14', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_14', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_14', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_14', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_14', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_15', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_15', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_15', 'Conv_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_15', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_15', 'Conv_2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_15', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_15', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_15', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_15', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_15', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_15', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_15', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_16', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_16', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_16', 'Conv_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_16', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_16', 'Conv_2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_16', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_16', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_16', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_16', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_16', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_16', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_16', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_17', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_17', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_17', 'Conv_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_17', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_17', 'Conv_2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_17', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_17', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_17', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_17', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_17', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_17', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_17', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_18', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_18', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_18', 'Conv_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_18', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_18', 'Conv_2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_18', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_18', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_18', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_18', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_18', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_18', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_18', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_19', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_19', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_19', 'Conv_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_19', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_19', 'Conv_2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_19', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_19', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_19', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_19', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_19', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_19', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_19', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_2', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_2', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_2', 'Conv_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_2', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_2', 'Conv_2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_2', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_2', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_2', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_2', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_2', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_2', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_2', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_20', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_20', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_20', 'Conv_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_20', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_20', 'Conv_2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_20', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_20', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_20', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_20', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_20', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_20', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_20', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_21', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_21', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_21', 'Conv_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_21', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_21', 'Conv_2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_21', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_21', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_21', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_21', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_21', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_21', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_21', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_3', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_3', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_3', 'Conv_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_3', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_3', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_3', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_3', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_3', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_3', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_3', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_4', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_4', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_4', 'Conv_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_4', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_4', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_4', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_4', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_4', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_4', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_4', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_5', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_5', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_5', 'Conv_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_5', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_5', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_5', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_5', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_5', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_5', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_5', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_6', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_6', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_6', 'Conv_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_6', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_6', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_6', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_6', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_6', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_6', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_6', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_7', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_7', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_7', 'Conv_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_7', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_7', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_7', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_7', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_7', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_7', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_7', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_8', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_8', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_8', 'Conv_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_8', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_8', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_8', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_8', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_8', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_8', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_8', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_9', 'Conv_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_9', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_9', 'Conv_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_9', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_9', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_9', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_9', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_9', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_9', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'ResnetBlock_9', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'TimeEmbedding_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "TimeEmbedding_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'TimeEmbedding_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "TimeEmbedding_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'TimeEmbedding_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "TimeEmbedding_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'TimeEmbedding_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "TimeEmbedding_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'UpSample_0', 'ConvTranspose_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UpSample_0", "key_type": 2}, {"key": "ConvTranspose_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'UpSample_0', 'ConvTranspose_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UpSample_0", "key_type": 2}, {"key": "ConvTranspose_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'UpSample_1', 'ConvTranspose_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UpSample_1", "key_type": 2}, {"key": "ConvTranspose_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'UpSample_1', 'ConvTranspose_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UpSample_1", "key_type": 2}, {"key": "ConvTranspose_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'UpSample_2', 'ConvTranspose_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UpSample_2", "key_type": 2}, {"key": "ConvTranspose_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'params', 'UpSample_2', 'ConvTranspose_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UpSample_2", "key_type": 2}, {"key": "ConvTranspose_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'count')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'DownSample_0', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "DownSample_0", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'DownSample_0', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "DownSample_0", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'DownSample_1', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "DownSample_1", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'DownSample_1', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "DownSample_1", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'DownSample_2', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "DownSample_2", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'DownSample_2', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "DownSample_2", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_0', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_0', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_1', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_1', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_2', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_2', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_3', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_3', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_4', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_4', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_5', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_5', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_0', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_0', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_0', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_0', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_0', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_0', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_0', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_0', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_1', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_1', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_1', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_1', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_1', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_1', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_1', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_1', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_1', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_1', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_10', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_10', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_10', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_10', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_10', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_10', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_10', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_10', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_10', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_10', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_10', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_10', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_11', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_11', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_11', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_11', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_11', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_11', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_11', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_11', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_11', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_11', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_11', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_11', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_12', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_12', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_12', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_12', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_12', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_12', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_12', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_12', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_12', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_12', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_12', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_12', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_13', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_13', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_13', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_13', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_13', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_13', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_13', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_13', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_13', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_13', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_13', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_13', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_14', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_14', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_14', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_14', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_14', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_14', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_14', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_14', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_14', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_14', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_14', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_14', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_15', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_15', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_15', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_15', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_15', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_15', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_15', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_15', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_15', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_15', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_15', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_15', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_16', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_16', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_16', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_16', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_16', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_16', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_16', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_16', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_16', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_16', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_16', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_16', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_17', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_17', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_17', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_17', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_17', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_17', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_17', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_17', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_17', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_17', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_17', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_17', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_18', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_18', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_18', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_18', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_18', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_18', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_18', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_18', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_18', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_18', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_18', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_18', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_19', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_19', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_19', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_19', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_19', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_19', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_19', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_19', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_19', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_19', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_19', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_19', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_2', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_2', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_2', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_2', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_2', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_2', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_2', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_2', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_2', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_2', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_2', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_2', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_20', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_20', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_20', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_20', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_20', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_20', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_20', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_20', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_20', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_20', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_20', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_20', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_21', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_21', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_21', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_21', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_21', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_21', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_21', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_21', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_21', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_21', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_21', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_21', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_3', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_3', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_3', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_3', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_3', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_3', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_3', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_3', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_3', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_3', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_4', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_4', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_4', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_4', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_4', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_4', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_4', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_4', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_4', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_4', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_5', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_5', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_5', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_5', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_5', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_5', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_5', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_5', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_5', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_5', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_6', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_6', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_6', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_6', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_6', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_6', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_6', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_6', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_6', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_6', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_7', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_7', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_7', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_7', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_7', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_7', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_7', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_7', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_7', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_7', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_8', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_8', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_8', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_8', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_8', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_8', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_8', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_8', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_8', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_8', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_9', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_9', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_9', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_9', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_9', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_9', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_9', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_9', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_9', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'ResnetBlock_9', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'TimeEmbedding_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "TimeEmbedding_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'TimeEmbedding_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "TimeEmbedding_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'TimeEmbedding_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "TimeEmbedding_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'TimeEmbedding_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "TimeEmbedding_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'UpSample_0', 'ConvTranspose_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UpSample_0", "key_type": 2}, {"key": "ConvTranspose_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'UpSample_0', 'ConvTranspose_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UpSample_0", "key_type": 2}, {"key": "ConvTranspose_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'UpSample_1', 'ConvTranspose_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UpSample_1", "key_type": 2}, {"key": "ConvTranspose_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'UpSample_1', 'ConvTranspose_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UpSample_1", "key_type": 2}, {"key": "ConvTranspose_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'UpSample_2', 'ConvTranspose_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UpSample_2", "key_type": 2}, {"key": "ConvTranspose_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'mu', 'params', 'UpSample_2', 'ConvTranspose_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UpSample_2", "key_type": 2}, {"key": "ConvTranspose_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'DownSample_0', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "DownSample_0", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'DownSample_0', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "DownSample_0", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'DownSample_1', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "DownSample_1", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'DownSample_1', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "DownSample_1", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'DownSample_2', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "DownSample_2", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'DownSample_2', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "DownSample_2", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_0', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_0', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_1', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_1', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_2', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_2', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_3', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_3', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_4', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_4', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_5', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_5', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_0', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_0', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_0', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_0', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_0', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_0', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_0', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_0', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_1', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_1', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_1', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_1', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_1', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_1', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_1', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_1', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_1', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_1', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_10', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_10', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_10', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_10', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_10', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_10', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_10', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_10', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_10', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_10', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_10', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_10', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_11', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_11', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_11', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_11', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_11', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_11', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_11', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_11', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_11', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_11', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_11', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_11', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_12', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_12', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_12', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_12', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_12', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_12', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_12', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_12', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_12', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_12', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_12', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_12', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_13', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_13', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_13', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_13', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_13', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_13', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_13', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_13', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_13', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_13', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_13', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_13', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_14', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_14', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_14', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_14', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_14', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_14', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_14', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_14', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_14', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_14', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_14', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_14', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_15', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_15', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_15', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_15', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_15', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_15', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_15', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_15', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_15', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_15', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_15', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_15', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_16', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_16', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_16', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_16', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_16', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_16', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_16', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_16', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_16', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_16', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_16', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_16', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_17', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_17', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_17', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_17', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_17', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_17', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_17', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_17', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_17', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_17', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_17', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_17', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_18', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_18', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_18', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_18', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_18', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_18', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_18', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_18', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_18', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_18', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_18', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_18', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_19', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_19', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_19', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_19', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_19', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_19', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_19', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_19', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_19', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_19', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_19', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_19', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_2', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_2', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_2', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_2', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_2', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_2', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_2', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_2', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_2', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_2', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_2', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_2', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_20', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_20', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_20', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_20', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_20', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_20', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_20', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_20', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_20', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_20', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_20', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_20', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_21', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_21', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_21', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_21', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_21', 'Conv_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_21', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_21', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_21', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_21', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_21', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_21', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_21', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_3', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_3', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_3', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_3', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_3', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_3', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_3', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_3', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_3', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_3', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_4', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_4', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_4', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_4', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_4', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_4', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_4', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_4', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_4', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_4', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_5', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_5', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_5', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_5', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_5', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_5', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_5', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_5', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_5', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_5', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_6', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_6', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_6', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_6', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_6', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_6', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_6', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_6', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_6', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_6', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_7', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_7', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_7', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_7', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_7', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_7', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_7', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_7', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_7', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_7', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_8', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_8', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_8', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_8', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_8', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_8', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_8', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_8', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_8', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_8', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_9', 'Conv_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_9', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_9', 'Conv_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_9', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_9', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_9', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_9', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_9', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_9', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'ResnetBlock_9', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'TimeEmbedding_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "TimeEmbedding_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'TimeEmbedding_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "TimeEmbedding_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'TimeEmbedding_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "TimeEmbedding_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'TimeEmbedding_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "TimeEmbedding_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'UpSample_0', 'ConvTranspose_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UpSample_0", "key_type": 2}, {"key": "ConvTranspose_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'UpSample_0', 'ConvTranspose_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UpSample_0", "key_type": 2}, {"key": "ConvTranspose_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'UpSample_1', 'ConvTranspose_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UpSample_1", "key_type": 2}, {"key": "ConvTranspose_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'UpSample_1', 'ConvTranspose_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UpSample_1", "key_type": 2}, {"key": "ConvTranspose_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'UpSample_2', 'ConvTranspose_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UpSample_2", "key_type": 2}, {"key": "ConvTranspose_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '0', 'nu', 'params', 'UpSample_2', 'ConvTranspose_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UpSample_2", "key_type": 2}, {"key": "ConvTranspose_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('opt_state', '1')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('rng',)": {"key_metadata": [{"key": "rng", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('train_metrics', '_reduction_counter', 'value')": {"key_metadata": [{"key": "train_metrics", "key_type": 2}, {"key": "_reduction_counter", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('train_metrics', 'step', 'total')": {"key_metadata": [{"key": "train_metrics", "key_type": 2}, {"key": "step", "key_type": 2}, {"key": "total", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('train_metrics', 'step', 'count')": {"key_metadata": [{"key": "train_metrics", "key_type": 2}, {"key": "step", "key_type": 2}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('train_metrics', 'epoch', 'total')": {"key_metadata": [{"key": "train_metrics", "key_type": 2}, {"key": "epoch", "key_type": 2}, {"key": "total", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('train_metrics', 'epoch', 'count')": {"key_metadata": [{"key": "train_metrics", "key_type": 2}, {"key": "epoch", "key_type": 2}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('model_metrics', '_reduction_counter', 'value')": {"key_metadata": [{"key": "model_metrics", "key_type": 2}, {"key": "_reduction_counter", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('model_metrics', 'loss', 'total')": {"key_metadata": [{"key": "model_metrics", "key_type": 2}, {"key": "loss", "key_type": 2}, {"key": "total", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('model_metrics', 'loss', 'count')": {"key_metadata": [{"key": "model_metrics", "key_type": 2}, {"key": "loss", "key_type": 2}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('train_dataset', 'epoch')": {"key_metadata": [{"key": "train_dataset", "key_type": 2}, {"key": "epoch", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('train_dataset', 'step')": {"key_metadata": [{"key": "train_dataset", "key_type": 2}, {"key": "step", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('train_dataset', 'rng')": {"key_metadata": [{"key": "train_dataset", "key_type": 2}, {"key": "rng", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('train_dataset', 'img')": {"key_metadata": [{"key": "train_dataset", "key_type": 2}, {"key": "img", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('val_dataset', 'epoch')": {"key_metadata": [{"key": "val_dataset", "key_type": 2}, {"key": "epoch", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('val_dataset', 'step')": {"key_metadata": [{"key": "val_dataset", "key_type": 2}, {"key": "step", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('val_dataset', 'rng')": {"key_metadata": [{"key": "val_dataset", "key_type": 2}, {"key": "rng", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('val_dataset', 'img')": {"key_metadata": [{"key": "val_dataset", "key_type": 2}, {"key": "img", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'Conv_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'DownSample_0', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "DownSample_0", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'DownSample_0', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "DownSample_0", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'DownSample_1', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "DownSample_1", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'DownSample_1', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "DownSample_1", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'DownSample_2', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "DownSample_2", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'DownSample_2', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "DownSample_2", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_0', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_0', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_0', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_1', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_1', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_1', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_2', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_2', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_2', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_3', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_3', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_3', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_4', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_4', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_4', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_5', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_5', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResidualAttentionBlock_5', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResidualAttentionBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_0', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_0', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_0', 'Conv_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_0', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_0', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_0', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_0', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_0', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_0", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_1', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_1', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_1', 'Conv_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_1', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_1', 'Dense_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_1', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_1', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_1', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_1', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_1', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_1", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_10', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_10', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_10', 'Conv_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_10', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_10', 'Conv_2', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_10', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_10', 'Dense_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_10', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_10', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_10', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_10', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_10', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_10", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_11', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_11', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_11', 'Conv_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_11', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_11', 'Conv_2', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_11', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_11', 'Dense_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_11', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_11', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_11', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_11', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_11', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_11", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_12', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_12', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_12', 'Conv_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_12', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_12', 'Conv_2', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_12', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_12', 'Dense_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_12', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_12', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_12', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_12', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_12', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_12", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_13', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_13', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_13', 'Conv_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_13', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_13', 'Conv_2', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_13', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_13', 'Dense_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_13', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_13', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_13', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_13', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_13', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_13", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_14', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_14', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_14', 'Conv_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_14', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_14', 'Conv_2', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_14', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_14', 'Dense_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_14', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_14', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_14', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_14', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_14', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_14", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_15', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_15', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_15', 'Conv_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_15', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_15', 'Conv_2', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_15', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_15', 'Dense_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_15', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_15', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_15', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_15', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_15', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_15", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_16', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_16', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_16', 'Conv_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_16', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_16', 'Conv_2', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_16', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_16', 'Dense_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_16', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_16', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_16', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_16', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_16', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_16", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_17', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_17', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_17', 'Conv_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_17', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_17', 'Conv_2', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_17', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_17', 'Dense_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_17', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_17', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_17', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_17', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_17', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_17", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_18', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_18', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_18', 'Conv_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_18', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_18', 'Conv_2', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_18', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_18', 'Dense_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_18', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_18', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_18', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_18', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_18', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_18", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_19', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_19', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_19', 'Conv_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_19', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_19', 'Conv_2', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_19', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_19', 'Dense_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_19', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_19', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_19', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_19', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_19', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_19", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_2', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_2', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_2', 'Conv_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_2', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_2', 'Conv_2', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_2', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_2', 'Dense_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_2', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_2', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_2', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_2', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_2', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_2", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_20', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_20', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_20', 'Conv_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_20', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_20', 'Conv_2', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_20', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_20', 'Dense_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_20', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_20', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_20', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_20', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_20', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_20", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_21', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_21', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_21', 'Conv_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_21', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_21', 'Conv_2', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_21', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_21', 'Dense_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_21', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_21', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_21', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_21', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_21', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_21", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_3', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_3', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_3', 'Conv_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_3', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_3', 'Dense_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_3', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_3', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_3', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_3', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_3', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_3", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_4', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_4', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_4', 'Conv_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_4', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_4', 'Dense_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_4', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_4', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_4', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_4', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_4', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_4", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_5', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_5', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_5', 'Conv_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_5', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_5', 'Dense_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_5', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_5', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_5', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_5', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_5', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_5", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_6', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_6', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_6', 'Conv_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_6', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_6', 'Dense_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_6', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_6', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_6', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_6', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_6', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_6", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_7', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_7', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_7', 'Conv_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_7', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_7', 'Dense_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_7', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_7', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_7', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_7', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_7', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_7", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_8', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_8', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_8', 'Conv_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_8', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_8', 'Dense_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_8', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_8', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_8', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_8', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_8', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_8", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_9', 'Conv_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_9', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_9', 'Conv_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_9', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_9', 'Dense_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_9', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_9', 'GroupNorm_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_9', 'GroupNorm_0', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "GroupNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_9', 'GroupNorm_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'ResnetBlock_9', 'GroupNorm_1', 'scale')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ResnetBlock_9", "key_type": 2}, {"key": "GroupNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'TimeEmbedding_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "TimeEmbedding_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'TimeEmbedding_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "TimeEmbedding_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'TimeEmbedding_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "TimeEmbedding_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'TimeEmbedding_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "TimeEmbedding_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'UpSample_0', 'ConvTranspose_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UpSample_0", "key_type": 2}, {"key": "ConvTranspose_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'UpSample_0', 'ConvTranspose_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UpSample_0", "key_type": 2}, {"key": "ConvTranspose_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'UpSample_1', 'ConvTranspose_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UpSample_1", "key_type": 2}, {"key": "ConvTranspose_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'UpSample_1', 'ConvTranspose_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UpSample_1", "key_type": 2}, {"key": "ConvTranspose_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'UpSample_2', 'ConvTranspose_0', 'bias')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UpSample_2", "key_type": 2}, {"key": "ConvTranspose_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('ema_params', 'params', 'UpSample_2', 'ConvTranspose_0', 'kernel')": {"key_metadata": [{"key": "ema_params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UpSample_2", "key_type": 2}, {"key": "ConvTranspose_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}}, "use_zarr3": false}