trouble deserializing a equinox tree leaf that has size [1,1]

[csafta]

Hello,
I am getting the following error message when trying to instantiate a model from a saved solution with one of the leaves having size [1,1]. Error message is copied below.
Thank you,
-Cosmin

File "/Users/csafta/opt/miniconda3/envs/jax_nonmetal/lib/python3.12/site-packages/equinox/_serialisation.py", line 333, in tree_deserialise_leaves
jtu.tree_map_with_path(_assert_same(array_impl_type), out, like, is_leaf=is_leaf)
File "/Users/csafta/opt/miniconda3/envs/jax_nonmetal/lib/python3.12/site-packages/jax/_src/tree_util.py", line 1200, in tree_map_with_path
return treedef.unflatten(f(*xs) for xs in zip(*all_keypath_leaves))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/csafta/opt/miniconda3/envs/jax_nonmetal/lib/python3.12/site-packages/jax/_src/tree_util.py", line 1200, in
return treedef.unflatten(f(*xs) for xs in zip(*all_keypath_leaves))
^^^^^^
File "/Users/csafta/opt/miniconda3/envs/jax_nonmetal/lib/python3.12/site-packages/equinox/_serialisation.py", line 178, in _assert_same_impl
raise RuntimeError(
RuntimeError: Deserialised leaf at path (GetAttrKey(name='_model'), GetAttrKey(name='layers'), SequenceKey(idx=0), GetAttrKey(name='weight')) has changed shape from (1, 1) in like to (1,) on disk.

[johannahaffner]

Hi Cosmin,

can you post a small snippet of code that reproduces the issue? (Ideally 20 lines or less.)