| {"tree_metadata": {"('buffer', 'proprio_obs')": {"key_metadata": [{"key": "buffer", "key_type": 2}, {"key": "proprio_obs", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [100000, 14]}}, "('buffer', 'privileged_obs')": {"key_metadata": [{"key": "buffer", "key_type": 2}, {"key": "privileged_obs", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [100000, 2]}}, "('buffer', 'vision_obs')": {"key_metadata": [{"key": "buffer", "key_type": 2}, {"key": "vision_obs", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [100000, 64, 64, 1]}}, "('buffer', 'action_chunks')": {"key_metadata": [{"key": "buffer", "key_type": 2}, {"key": "action_chunks", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [100000, 16, 7]}}, "('buffer', 'ptr')": {"key_metadata": [{"key": "buffer", "key_type": 2}, {"key": "ptr", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('buffer', 'size')": {"key_metadata": [{"key": "buffer", "key_type": 2}, {"key": "size", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('student', 'params', 'Dense_0', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('student', 'params', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 2]}}, "('student', 'params', 'PhysicsHead_0', 'LowRankDense_0', 'base', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "PhysicsHead_0", "key_type": 2}, {"key": "LowRankDense_0", "key_type": 2}, {"key": "base", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('student', 'params', 'PhysicsHead_0', 'LowRankDense_0', 'base', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "PhysicsHead_0", "key_type": 2}, {"key": "LowRankDense_0", "key_type": 2}, {"key": "base", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('student', 'params', 'PhysicsHead_0', 'LowRankDense_0', 'lora_A', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "PhysicsHead_0", "key_type": 2}, {"key": "LowRankDense_0", "key_type": 2}, {"key": "lora_A", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 8]}}, "('student', 'params', 'PhysicsHead_0', 'LowRankDense_0', 'lora_B', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "PhysicsHead_0", "key_type": 2}, {"key": "LowRankDense_0", "key_type": 2}, {"key": "lora_B", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('student', 'params', 'PhysicsHead_0', 'LowRankDense_1', 'base', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "PhysicsHead_0", "key_type": 2}, {"key": "LowRankDense_1", "key_type": 2}, {"key": "base", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('student', 'params', 'PhysicsHead_0', 'LowRankDense_1', 'base', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "PhysicsHead_0", "key_type": 2}, {"key": "LowRankDense_1", "key_type": 2}, {"key": "base", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 2]}}, "('student', 'params', 'PhysicsHead_0', 'LowRankDense_1', 'lora_A', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "PhysicsHead_0", "key_type": 2}, {"key": "LowRankDense_1", "key_type": 2}, {"key": "lora_A", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 8]}}, "('student', 'params', 'PhysicsHead_0', 'LowRankDense_1', 'lora_B', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "PhysicsHead_0", "key_type": 2}, {"key": "LowRankDense_1", "key_type": 2}, {"key": "lora_B", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 2]}}, "('student', 'params', 'UNet1D_0', 'Conv_0', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UNet1D_0", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('student', 'params', 'UNet1D_0', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UNet1D_0", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 263, 256]}}, "('student', 'params', 'UNet1D_0', 'Conv_1', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UNet1D_0", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('student', 'params', 'UNet1D_0', 'Conv_1', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UNet1D_0", "key_type": 2}, {"key": "Conv_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('student', 'params', 'UNet1D_0', 'Conv_2', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UNet1D_0", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('student', 'params', 'UNet1D_0', 'Conv_2', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UNet1D_0", "key_type": 2}, {"key": "Conv_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 7]}}, "('student', 'params', 'UNet1D_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UNet1D_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('student', 'params', 'UNet1D_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UNet1D_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 128]}}, "('student', 'params', 'UNet1D_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UNet1D_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('student', 'params', 'UNet1D_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "UNet1D_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [270, 128]}}, "('student', 'params', 'VisionEncoder_0', 'Conv_0', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('student', 'params', 'VisionEncoder_0', 'Conv_0', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "Conv_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 8, 1, 128]}}, "('student', 'params', 'VisionEncoder_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('student', 'params', 'VisionEncoder_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_0', 'LayerNorm_0', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_0", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_0', 'LayerNorm_0', 'scale')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_0", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_0', 'LayerNorm_1', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_0", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_0', 'LayerNorm_1', 'scale')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_0", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_0', 'MlpBlock_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_0", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_0', 'MlpBlock_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_0", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_0', 'MlpBlock_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_0", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_0', 'MlpBlock_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_0", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 128]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_0', 'SelfAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_0", "key_type": 2}, {"key": "SelfAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_0', 'SelfAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_0", "key_type": 2}, {"key": "SelfAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 4, 32]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_0', 'SelfAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_0", "key_type": 2}, {"key": "SelfAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_0', 'SelfAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_0", "key_type": 2}, {"key": "SelfAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32, 128]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_0', 'SelfAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_0", "key_type": 2}, {"key": "SelfAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_0', 'SelfAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_0", "key_type": 2}, {"key": "SelfAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 4, 32]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_0', 'SelfAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_0", "key_type": 2}, {"key": "SelfAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_0', 'SelfAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_0", "key_type": 2}, {"key": "SelfAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 4, 32]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_1', 'LayerNorm_0', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_1", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_1', 'LayerNorm_0', 'scale')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_1", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_1', 'LayerNorm_1', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_1", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_1', 'LayerNorm_1', 'scale')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_1", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_1', 'MlpBlock_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_1", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_1', 'MlpBlock_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_1", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_1', 'MlpBlock_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_1", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_1', 'MlpBlock_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_1", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 128]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_1', 'SelfAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_1", "key_type": 2}, {"key": "SelfAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_1', 'SelfAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_1", "key_type": 2}, {"key": "SelfAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 4, 32]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_1', 'SelfAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_1", "key_type": 2}, {"key": "SelfAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_1', 'SelfAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_1", "key_type": 2}, {"key": "SelfAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32, 128]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_1', 'SelfAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_1", "key_type": 2}, {"key": "SelfAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_1', 'SelfAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_1", "key_type": 2}, {"key": "SelfAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 4, 32]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_1', 'SelfAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_1", "key_type": 2}, {"key": "SelfAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('student', 'params', 'VisionEncoder_0', 'EncoderBlock_1', 'SelfAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "EncoderBlock_1", "key_type": 2}, {"key": "SelfAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 4, 32]}}, "('student', 'params', 'VisionEncoder_0', 'LayerNorm_0', 'bias')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('student', 'params', 'VisionEncoder_0', 'LayerNorm_0', 'scale')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('student', 'params', 'VisionEncoder_0', 'cls_token')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "cls_token", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128]}}, "('student', 'params', 'VisionEncoder_0', 'pos_embedding')": {"key_metadata": [{"key": "student", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "VisionEncoder_0", "key_type": 2}, {"key": "pos_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 65, 128]}}, "('teacher', 'params', 'Dense_0', 'bias')": {"key_metadata": [{"key": "teacher", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('teacher', 'params', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "teacher", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 256]}}, "('teacher', 'params', 'Dense_1', 'bias')": {"key_metadata": [{"key": "teacher", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('teacher', 'params', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "teacher", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('teacher', 'params', 'Dense_2', 'bias')": {"key_metadata": [{"key": "teacher", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('teacher', 'params', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "teacher", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 7]}}, "('teacher', 'params', 'Dense_3', 'bias')": {"key_metadata": [{"key": "teacher", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('teacher', 'params', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "teacher", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 1]}}, "('teacher', 'params', 'log_std')": {"key_metadata": [{"key": "teacher", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "log_std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} |