How to return an object in a function while calculating its gradient/Jacobian?

[GaoyuanWu]

Hi all.
Say I have the following definitions:

#A class
class c1():
  def __init__(self,x):
    self.x = x

#A function
def f(x):
  y = x**2
  obj = c1(x) #an object
  return y,obj

And I wish to calculate $\frac{dy}{dx}$ using jax.grad and get $y$ and obj at the same time.

grad_y_x = jax.grad(f,has_aux=True)
res = grad_y_x(1)

This will fail as "<class 'c1'> is not a valid JAX type. How do I get both the f(x), which contains an object return, and $\frac{dy}{dx}$ at the same time?
Thanks so much!!

[jakevdp]

In order to do this you need to register your class c1 as a pytree. You can find more information here: https://jax.readthedocs.io/en/latest/pytrees.html#extending-pytrees

For example:

import jax
import jax.numpy as jnp
from jax.tree_util import register_pytree_node

#A class
class c1():
  def __init__(self, x):
    self.x = x

def c1_flatten(obj):
  aux_data = None
  children = (obj.x,)
  return children, aux_data

def c1_unflatten(aux_data, children):
  return c1(*children)

register_pytree_node(c1, c1_flatten, c1_unflatten)

#A function
def f(x):
  y = x**2
  obj = c1(x) #an object
  return y,obj

jax.grad(f, has_aux=True)(0.0)
# (DeviceArray(0., dtype=float32), <__main__.c1 at 0x7fcdf7bd1fd0>)