compatibility issue with submitit, which depends on cloudpickle

[eringrant]

Related to #198 as an AttributeError with ...has no attribute 'index_variadic', but occurs in __eq__ and seemingly only for PRNGKeyArray annotations:

Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/usr/local/lib/python3.12/site-packages/submitit/core/_submit.py", line 11, in <module>
    submitit_main()
  File "/usr/local/lib/python3.12/site-packages/submitit/core/submission.py", line 76, in submitit_main
    process_job(args.folder)
  File "/usr/local/lib/python3.12/site-packages/submitit/core/submission.py", line 69, in process_job
    raise error
  File "/usr/local/lib/python3.12/site-packages/submitit/core/submission.py", line 52, in process_job
    delayed = utils.DelayedSubmission.load(paths.submitted_pickle)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/submitit/core/utils.py", line 153, in load
    obj = pickle_load(filepath)
          ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/submitit/core/utils.py", line 232, in pickle_load
    return pickle.load(ifile)
           ^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/cloudpickle/cloudpickle.py", line 549, in _make_skeleton_class
    return _lookup_class_or_track(class_tracker_id, skeleton_class)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/cloudpickle/cloudpickle.py", line 124, in _lookup_class_or_track
    _DYNAMIC_CLASS_TRACKER_BY_CLASS[class_def] = class_tracker_id
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^
  File "/usr/local/lib/python3.12/weakref.py", line 428, in __setitem__
    self.data[ref(key, self._remove)] = value
    ~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/jaxtyping/_array_types.py", line 337, in __eq__
    return cls._get_props() == other._get_props()
                               ^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/jaxtyping/_array_types.py", line 325, in _get_props
    cls.index_variadic,
    ^^^^^^^^^^^^^^^^^^
AttributeError: type object 'UInt32[Array, '2']' has no attribute 'index_variadic'

I don't understand how to generalize the solution in #237 to this case, but I've got a MWE:

import jax
import submitit

from jaxtyping import Array
from jaxtyping import PRNGKeyArray


#def noise(key) -> Array:  # removing the key annotation avoids the error
def noise(key: PRNGKeyArray) -> Array:
  return jax.random.normal(key)


def main():
  executor = submitit.AutoExecutor(folder="/tmp/test_submitit", cluster="local")
  job = executor.submit(noise, jax.random.PRNGKey(0))
  output = job.results()
  print(output)


if __name__ == "__main__":
  main()

My package versions are:

python -m pip list | grep "jaxtyping\|equinox\|cloudpickle\|submitit"
cloudpickle               3.1.0
equinox                   0.11.8
jaxtyping                 0.2.34
submitit                  1.5.2

[patrick-kidger]

Thanks for the report! Looks like a variation on the same thing as before, cloudpickle is trying to hash/eq partially-constructed objects.

I've just opened #262 to hopefully solve this once and for all. I don't think us supporting 'nicer' hash/eq is as important as having deserialization work smoothly.

[eringrant]

Confirming #262 avoids the issue for me. Thanks so much for the quick and pragmatic fix!