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

Open
#343 4 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
45/100
Issue type
Bug
Clarity
Mostly clear
Activity status
Stale
Tech stack
python
Domain
tooling

Research direction

Start by reproducing the example with cloudpickle, jaxtyping, and typeguard, then inspect the decorator entry point in jaxtyping/_decorator.py around line 71. Compare pickling the decorated function with the unwrapped function; done means the decorated function can be serialized without the weakref.ReferenceType error while preserving the reported behavior.

Written by the indexing model from the issue text.

Description

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

Dominant language
Python
Stars
1.9k
Forks
96
No merged PRs in 30d

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from patrick-kidger/jaxtyping

All issues in patrick-kidger/jaxtyping

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.