cloudpipe / cloudpipe/cloudpickle

jax jitted functions cloudpickled work but include some error messages

Open
#537 4 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
1.9k
Forks
195
Avg merge
1d 10h
Merged PRs (30d)
1

Description

# problems

cloudpickle works for jax.jit functions but a visual inspection of the cloudpickle contents shows there's a lurking error message

# challenges

not sure if this belongs in cloudpickle or jax

is this my bad? I was hopeful we could just use the string jaxpr in utf8, it's more human readable, but I don't know how to regenerate a PjitFunction from a jaxpr

# opportunities

a fix could reduce the size of cloudpickled jax.jit functions

```py
def test_jax_cloudpickle():
def jnp_func(x):
return jax.numpy.sin(jax.numpy.cos(x))

jitted1 = jax.jit(jnp_func)
del jnp_func # this to ensure jitted2 can't cheat by recompiling jnp_func within a session
assert "jnp_func" not in locals(), "failed to remove jnp_func"
jitted1_buf = cloudpickle.dumps(jitted1)
rprint(jitted1_buf)
jitted2 = cloudpickle.loads(jitted1_buf)
assert jitted1(0.3) == jitted2(0.3), "weird"
assert b"TRACEBACK" not in jitted1_buf, "error message in cloudpickle of jax.jit"

test_jax_cloudpickle()
```

Could `JAX_TRACEBACK_FILTERING=` be greppable ?

![image](https://github.com/cloudpipe/cloudpickle/assets/24532336/c70faea5-6c07-45e6-a22c-04b73245b7e5)

thank you for making cloudpickle

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.