How should the Tracers be handled while using pytrees to register a custom python container?

[trickarcher]

Hi,

We are trying to use pytrees, to register Awkward Arrays(which is a tree-like python container, and fits perfectly in the use-case definition of pytrees). to allow JAX to handle element wise ops on it. Just for some context, an Awkward Array can be broken down into a (form, length, children) where, (form, length) correspond to the aux_data of the pytrees and children are a list of linear buffers. We wrote our own version of special_unflatten and special_flatten for one type of Awkward Arrays called the ak.layout.NumpyArray which is essentially same as a numpy.ndarray.

While calling the jvp function, we recieve JVPTracers. From my limited knowledge on JAX, it has a primal and a tangent. But to form an Awkward Array we need a concrete linear buffer, which the JVPTracer does provide, but it gives two of them. How should we be handling this tracer? For a simple example, if we catch the JVPTracer object using isinstance and return it as-is, we do get the correct gradients. However, from the one example on the documentation it seems that these Tracers need to be somehow in the form of the python container. If so, how could we convert something like a Tracer to a numpy-array like object without losing the primal and the tangent information?

EDIT: Here's a basic code example, taking in a wrapper NumpyArray class and attempting a jvp pass on it.

import jax
from jax.tree_util import tree_flatten, tree_unflatten, register_pytree_node, all_leaves
import numpy as np
jax.config.update("jax_platform_name", "cpu")
jax.config.update("jax_enable_x64", True)
class NumpyArray(object):
    def __init__(self, array):
        self.aux_data = (array.dtype, len(array))
        self.data = array
    def __getitem__(self,key):
        return self.data[key]
def special_flatten(v):
    dtype, length = v.aux_data
    children = [jax.numpy.asarray(v.data)]
    return (children, v.aux_data)
def special_unflatten(aux_data, children):
    dtype, length = aux_data
    children = list(children)
    # No way to form an Numpy Array out of a None buffer with a given aux_data so just return the aux_data 
    if (children[0] == None):
      return (dtype, length)
    # If there's a sentinel object recieved return a placeholder buffer of zeros
    if type(children[0]) == object:
        children[0] = np.zeros(length)
    # Catch the JVPTracer here and return it as is. The gradients are correct.
    # If I return the primal as an NumpyArray I get the wrong gradients but the correct 
    # value. 
    # Returning the primal and tangent as a tuple returns a 
    # TypeError: unsupported operand type(s) for ** or pow(): 'NumpyArray' and 'int'
    if isinstance(children[0], jax.interpreters.ad.JVPTracer):
        return children[0] # JVP gives the correct results here, (DeviceArray(5., dtype=float64), DeviceArray(4., dtype=float64))
        # return NumpyArray(np.array(children[0].primal)) # JVP gives the wrong gradients here, since we lose the tangent information: (5.0, array(0.))
        # return (NumpyArray(np.array(children[0].primal)), NumpyArray(np.array(children[0].tangent))) # Returns the TypeError
    return NumpyArray(children[0])
register_pytree_node(
    NumpyArray,
    special_flatten,  
    special_unflatten
)
def func(x):
    return x[0] ** 2 + x[1] ** 2
tree = NumpyArray(np.array([1., 2., 3., 4.]))
basis = NumpyArray(np.array([0., 1., 0., 0.]))
print(jax.jvp(func, (tree,), (basis,)))

[jakevdp]

I'm having trouble understanding from your example what your question is: in your implementation, x[0] and x[1] return normal numpy scalars, and when you add these together it returns a normal numpy scalar. The output of the function is exactly the same as if you removed the custom array wrapper and changed it to the following:

def func(x):
    return x[0] ** 2 + x[1] ** 2
tree = np.array([1., 2., 3., 4.])
basis = np.array([0., 1., 0., 0.])
print(jax.jvp(func, (tree,), (basis,)))

and thus it seems to be working as expected. Was that your understanding?

[trickarcher]

Hi! Thanks for the quick follow-up. The custom array wrapper was an example to demonstrate that when we perform any function transformation, we get these intermediate tracer objects in the special_unflatten function, which can't be expressed in terms of the custom class. In the code example, the commented out return statements are the problematic ones which cause wrong gradients or TypeError (According to the notebook example, the tracers should have been expressed in the form of the RegisteredSpecial class). My question was initially that, to register more complex types in JAX, would we have to add some additional functionality to handles these JAX-specific Tracers, or is there any other alternative solution to this.

[jakevdp]

I see. I would think the correct approach would involve wrapping the tracer itself. For example, if you change special_unflatten to the following, the code appears to work correctly:

def special_unflatten(aux_data, children):
    dtype, length = aux_data
    return NumpyArray(children[0])

[trickarcher]

Ah, that's my bad. It was probably a bad example to begin with. So wrapping the Tracer with this "fake" class doesn't complain because the constructor doesn't need anything specific(device_buffer) to form it. In our use-case, we have ak.from_jax() which would convert normal jax.numpy.array to an Awkward Array. So this naturally fails in our case. I guess, the bigger question is that each of the Tracers have different attributes like the BatchedTracer has a val attribute wheras the JVPTracer has the primal and tangent attribute. So in my understanding we would have to implement special classes which closely mirror these Tracers but have the buffers like val, primal and tangent, as an Awkward Array. And all these special classes can then become separate pytrees. Would this be the only way to go about it?

[jakevdp]

At the level you're working at, you should not need to worry about what flavor of tracer you're dealing with – just operate on it as if it is a DeviceArray.

If that does not work for some reason (i.e. you need to branch based on the contents of the array, or convert to a regular numpy array) then your application is likely incompatible (or not easily compatible) with JAX transforms.

If you can show an example of such an incompatibility, we might be able to offer more pointed advice.