make custom_transforms handle pytrees, add api.defvjp

[mattjj]

Before this PR, using custom_transforms (from #419) with functions that took tuples/lists/dicts (i.e. pytrees) as arguments would silently fail. Moreover, defining custom rules that took as input or returned as output any pytree values wouldn't work. Finally, custom_transforms functions didn't automatically support vmap.

This PR fixes all that!

This PR also improves the API by adding defjvp and defvjp (cf. #636) functions (and their variants) to api.py, which can be called directly on a custom_transforms function rather than having to dig out a reference to the underlying primitive. There's also a custom_gradient convenience wrapper, which works like tf.custom_gradient.

This API will likely be the main mechanism by which users will define custom gradients.

TODO:

• add docstrings

In the future we should add a tutorial explanation of how to do this in an .md file, like in Autograd's tutorial, at which point we can finally close #116.

It would also be good to add more tests, especially for error messages.

[fehiepsi]

Hi @mattjj , will this support batching rule for transforms with 2 arguments? Currently, we have to use the following pattern for a custom transform with 2 arguments

def f_batching_rule(batched_args, batch_dims):
    x, y = batched_args
    bx, by = batch_dims
    # promote shapes
    sx, sy = np.shape(x), np.shape(y)
    nx = len(sx) + int(bx is None)
    ny = len(sy) + int(by is None)
    nd = max(nx, ny)
    x = np.reshape(x, (1,) * (nd - len(sx)) + sx)
    y = np.reshape(y, (1,) * (nd - len(sy)) + sy)
    # correct bx, by due to promoting
    bx = bx + nd - len(sx) if bx is not None else nd - len(sx) - 1
    by = by + nd - len(sy) if by is not None else nd - len(sy) - 1
    # move bx, by to front
    x = batching.move_dim_to_front(x, bx)
    y = batching.move_dim_to_front(y, by)
    return f(x, y), 0

Basically, the code promote x,y to have the same ndim by appending singleton dimensions to the front, then correct vmap dims due to this promoting. Finally, it moves vmap dims to the front.

[mattjj]

It should handle arbitrary numbers of arguments according to standard vmap semantics. Do you have a test case I can try?

[fehiepsi]

Sure, for example,

import jax.numpy as np
from jax import custom_transforms, vmap

@custom_transforms
def xlogy(x, y):
    return x * np.log(y)

vmap(xlogy)(np.ones(3), np.ones((3, 2)))

[mattjj]

That seems to work on this branch: it produces a zeros array of shape (3, 2). Is that the result you expect?

[fehiepsi]

Yup, that is the expected result. It is great to see that vmap is working for these custom transforms now.

[skye]

Still reviewing, gonna send these out in case you wanna pipeline reviewing + responding...

[skye]

Finished reviewing the original PR, gonna look at changes now...

[skye]

Maybe add a test for a function that takes multiple and/or pytree arguments? (Sorry if I missed it below)

[skye]

(This can be a follow-up change, but I think is worth doing, at the very least to make sure we have a runnable example of multiple args.)

[mattjj]

Will do in a follow-up; left a todo.

[mattjj]

@skye convinced me she was right about not having api.defjvp2/api.defvjp2 on API minimalism grounds (rather than readability grounds). I updated the PR to only have a defjvp and defvjp, though both act like the "2" variants in that they take ans as an argument (which can be ignored).

Currently that means that the api.py functions are slightly off from the corresponding ad.py versions, but we'll remedy that in a follow-up.

EDIT: one other thing to mention: it's easy to add convenience variants if we decide we want them, but it's hard to remove them once added.