Pickling Jaxprs

[gokul-uf]

Hi folks,
I'm generating jaxpr for the forward pass of a LLM. I'm using the generated jaxpr for some downstream processing.
Generating jaxpr takes ~10 mins and I'd like to save / cache this result to disk to save this time.

But, when I try to do a pkl.dump(jaxpr...), I get the following error:

AttributeError: Can't pickle local object '_effect_free_abstract_eval.<locals>.abstract_eval_'

Is there any way around this?
Alternatively, what's the recommended way to pickle jaxprs?

Thanks!

[jakevdp]

No, unfortunately there's no way in JAX's current release to pickle jaxprs. The issue is that jaxprs are built on JAX primitives, each of which has associated impl and abstract_eval rules, many of which are defined using local/anonymous functions that are not compatible with pickle. Because of that, making jaxprs fully compatible with pickle would require some level of rewriting of JAX's internal code.

[gokul-uf]

Thanks @jakevdp , that's unfortunate.

Are there any other serialization libraries that might work instead?

[jakevdp]

I don't know of any way to serialize jaxprs although the experimental persistent compilation cache might help depending on what you're trying to do (it's only implemented on TPU at the moment, I think)

[patrick-kidger]

It might be possible to use something like dill instead of pickle.

[EelcoHoogendoorn]

See this issue here for a related dicussion