Different PyTree for Primal and Tangent space

[llandsmeer]

Hi,

Say I have a PyTree that represents some data structure.

In some cases, it might be beneficial to have a different data structure in tangent space. For example, Primal might be a more complex buffer structure, and Tangent a lossy emulation because we want to save memory by storing a reduced representation of the gradient

As a requirement, that would mean that the only way to interact with such datastructure would be functions with custom_jvp's

As a very simplified example:

class Primal(typing.NamedTuple):
    x: float

class Tangent(typing.NamedTuple):
    sign: bool
    norm: float
    def add(self, n):
        # some random equation
        return Tangent(n>0, self.norm + jnp.abs(n))

A potential way to construct such object could be

## init function
@jax.custom_jvp
def init(dummy) -> Primal:
    return Primal(0.)

@init.defjvp
def init_jvp(primals, tangents):
    return Primal(0.), Tangent(0., 0.)

## add function
@jax.custom_jvp
def add(p: Primal, n: float) -> Primal:
    return Primal(p.x + n)

@add.defjvp
def add_jvp(primals, tangents):
    p: Primal
    t: Tangent # Wrong of course, t: Primal
    p, n = primals
    t, n_dot = tangents
    return add(p, n), t.add(n_dot)

This fails with

TypeError: Custom JVP rule init_jvp for function init must produce primal and tangent outputs with equal container (pytree) structures, but got PyTreeDef(CustomNode(namedtuple[Primal], [])) and PyTreeDef(CustomNode(namedtuple[Tangent], [, *])) respectively.

Thus, this it seems this is not supported in current jax. The add() function doesn't work as t is a Primal and not a Tangent.

Is there some way to achieve this in current jax or is this maybe planned in the future?

Thanks,
Lennart

[dfm]

Great question! I don't have all the relevant context here so perhaps @mattjj, @froystig, or @dougalm will chime in with more details, but the short answer is that the general thing that you're asking about isn't currently supported in JAX, although it's always been something we hope to support one day.

The more specific side point about using different dtypes for the primal and tangent is approximately supported as discussed in this issue comment, but it's a pretty experimental feature, and hopefully we can eventually support the more general use case!

[llandsmeer]

Thanks a lot for answering. The dtype thing I was not aware of, thank you for pointing that out! IIUC the API is a bit too rough for our current use case, but that's definitely something useful for another project.