Cloudpickle interaction: TypeError: cannot pickle 'weakref.ReferenceType' object

[ArthurConmy]

This code:

# pip install cloudpickle jaxtyping typeguard
import cloudpickle
from jaxtyping import jaxtyped
import typeguard

def typed(function):
  return jaxtyped(function, typechecker=typeguard.typechecked)

@typed
def f():
  return 1

def unwrapped_f():
  return 2

# This will succeed
pickled_unwrapped = cloudpickle.dumps(unwrapped_f)
print("Successfully pickled unwrapped_f")

# This will fail
try:
  pickled_f = cloudpickle.dumps(f)
except Exception as e:
  print(f"\nFailed to pickle decorated function f:\n{type(e).__name__}: {e}")

Prints out this:

[...]/lib/python3.13/site-packages/jaxtyping/_decorator.py:71: InstrumentationWarning: instrumentor did not find the target function -- not typechecking __main__.f
  return typechecker(fn)
Successfully pickled unwrapped_f

Failed to pickle decorated function f:
TypeError: cannot pickle 'weakref.ReferenceType' object

Showing that jaxtyping doesn't interact well with cloudpickle :(

CC @JoshEngels @jkramar

[ArthurConmy]

That was with Typeguard V4 but it also happens with earlier versions:

(gchat_download_venv) conmy@conmy2:~/Documents$ pip uninstall typeguard
Found existing installation: typeguard 4.4.4
Uninstalling typeguard-4.4.4:
  Would remove:
    [...]gchat_download_venv/lib/python3.13/site-packages/typeguard-4.4.4.dist-info/*
    [...]gchat_download_venv/lib/python3.13/site-packages/typeguard/*
Proceed (Y/n)? Y
  Successfully uninstalled typeguard-4.4.4
(gchat_download_venv) conmy@conmy2:~/Documents$ pip install typeguard==2.13.3
Collecting typeguard==2.13.3
  Downloading typeguard-2.13.3-py3-none-any.whl.metadata (3.6 kB)
Downloading typeguard-2.13.3-py3-none-any.whl (17 kB)
Installing collected packages: typeguard
Successfully installed typeguard-2.13.3
(gchat_download_venv) conmy@conmy2:~/Documents$ python3 [...]jt.py
 
Successfully pickled unwrapped_f

Failed to pickle decorated function f:
TypeError: cannot pickle 'weakref.ReferenceType' object
(gchat_download_venv) conmy@conmy2:~/Documents$ 

[johannahaffner]

This discussion is probably relevant for you: https://github.com/patrick-kidger/jaxtyping/issues/332.

[patrick-kidger]

Yup, this is a topic I'd be happy to take a PR on :) The root cause is clear to me, but I'm not immediately sure how to tackle it.

[ArthurConmy]

Gemini Deep Research made something: https://github.com/patrick-kidger/jaxtyping/compare/main...ArthurConmy:jaxtyping:deep?expand=1 in typical style, with excess changes...

EDIT: It also fails the test of adding back runtime typechecking when loading the cloudpickle, sad