How to make JVPTracer JSON serializable?

[tylerflex]

I am trying to save an object containing JVPTracer objects to json format (or more generally, some string that I can load back later into the original object). Using pickle is not an option for my application due to security concerns.

Is there a way to serialize objects containing VJPTracers using jax?

Here is a minimally reproducible example. Thanks a lot for your help! (jax version 0.4.7)

import jax
import json

def f(d: dict) -> float:
    json_string = json.dumps(d)
    # do something with json_string
    return d['key'] * 2

g = jax.grad(f)

d0 = dict(key=4.0)

f(d0)
g(d0)

which results in

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[9], line 13
     10 d0 = dict(key=4.0)
     12 f(d0)
---> 13 g(d0)

    [... skipping hidden 10 frame]

Cell In[9], line 5, in f(d)
      4 def f(d: dict) -> float:
----> 5     json.dumps(d)
      6     return d['key'] * 2

File ~/.pyenv/versions/3.10.9/lib/python3.10/json/__init__.py:231, in dumps(obj, skipkeys, ensure_ascii, check_circular, allow_nan, cls, indent, separators, default, sort_keys, **kw)
    226 # cached encoder
    227 if (not skipkeys and ensure_ascii and
    228     check_circular and allow_nan and
    229     cls is None and indent is None and separators is None and
    230     default is None and not sort_keys and not kw):
--> 231     return _default_encoder.encode(obj)
    232 if cls is None:
    233     cls = JSONEncoder

File ~/.pyenv/versions/3.10.9/lib/python3.10/json/encoder.py:199, in JSONEncoder.encode(self, o)
    195         return encode_basestring(o)
    196 # This doesn't pass the iterator directly to ''.join() because the
    197 # exceptions aren't as detailed.  The list call should be roughly
    198 # equivalent to the PySequence_Fast that ''.join() would do.
--> 199 chunks = self.iterencode(o, _one_shot=True)
    200 if not isinstance(chunks, (list, tuple)):
    201     chunks = list(chunks)

File ~/.pyenv/versions/3.10.9/lib/python3.10/json/encoder.py:257, in JSONEncoder.iterencode(self, o, _one_shot)
    252 else:
    253     _iterencode = _make_iterencode(
    254         markers, self.default, _encoder, self.indent, floatstr,
    255         self.key_separator, self.item_separator, self.sort_keys,
    256         self.skipkeys, _one_shot)
--> 257 return _iterencode(o, 0)

File ~/.pyenv/versions/3.10.9/lib/python3.10/json/encoder.py:179, in JSONEncoder.default(self, o)
    160 def default(self, o):
    161     """Implement this method in a subclass such that it returns
    162     a serializable object for ``o``, or calls the base implementation
    163     (to raise a ``TypeError``).
   (...)
    177 
    178     """
--> 179     raise TypeError(f'Object of type {o.__class__.__name__} '
    180                     f'is not JSON serializable')

TypeError: Object of type JVPTracer is not JSON serializable

[jakevdp]

Hi - thanks for the question! I think you'll find the answer you're looking for in the FAQ: https://jax.readthedocs.io/en/latest/faq.html#how-can-i-convert-a-jax-tracer-to-a-numpy-array

I suspect you'll want to use jax.experimental.io_callback, which unfortunately is not well-documented yet, but I spent the morning working on adding that documentation, so it should be soon!