Elementwise grad

[mgostIH]

Hello, I have a function of signature (N,) -> (N,) and I'd like to get the gradient with respect to each element individually, not the jacobian matrix.
Usually one can start from a function f: R -> R and use vmap(grad(f)), but it's unclear what to do for functions that are already defined to work on vectors or even taking other shapes (with input/output being the same shape)

[leakec]

I opened a similar issue a few years back here. Is the egrad function that @mattjj suggested what you're looking for?

[mgostIH]

Excellent, works exactly as expected and is easy to understand, thanks!

A bit off topic, do you also happen to know what's the best way to turn a multidimensional array with a single element into a scalar? Right now I'm using .sum()

[jakevdp]

I think the most direct way to do that is x = arr.reshape(())

[YouJiacheng]

BTW, all ones jvp maybe faster than vjp.

[leakec]

@YouJiacheng In the final version I did end up using jvp rather than vjp. See comment below for more details.

[leakec]

Sorry I didn't throw this on the initial post as well. Not sure if you will need this, but the egrad function can be extended to encompass general pytrees by using @partial(partial, tree_multimap). See egrad_robust on line 131 here. @YouJiacheng In that version I did switch to jvp rather than vjp, can't remember why not, but it may have been because of speed.

[leakec]

Decided to just throw the code here so you don't have to dig around in the larger document. In this chunk of code, onp references the original numpy, i.e., import numpy as onp, and np references JAX's numpy, i.e., import jax.numpy as np.

I used onp ones and zeros here because I use egradRobust in a way later on that requires the tans to be numpy arrays rather than JAX numpy arrays. However, for most applications I think you could just use regular JAX numpy arrays instead.

@partial(partial, tree_multimap)
def onesRobust(val):
    """Returns ones_like val, but can handle arrays and dictionaries.
    Parameters
    ----------
    val : pytree
    Returns
    -------
    ones_like_val : pytree
        Pytree with the same structure as val with all elements equal to one.
    """
    return onp.ones(val.shape, dtype=val.dtype)


@partial(partial, tree_multimap)
def zerosRobust(val):
    """Returns zeros_like val, but can handle arrays and dictionaries.
    Parameters
    ----------
    val : pytree
    Returns
    -------
    zeros_like_val : pytree
        Pytree with the same structure as val with all elements equal to zero.
    """
    return onp.zeros(val.shape, dtype=val.dtype)


def egradRobust(g, j=0):
    """This function mimics egrad from autograd, but can also handle dictionaries.
    Parameters
    ----------
    g : function
        Function to take the derivative of.
    j : integer, optional
        Parameter with which to take the derivative with respect to. (Default value = 0)
    Returns
    -------
    wrapped : function
        Derivative function
    """
    if g.__qualname__ == "jit.<locals>.f_jitted":
        g = g.__wrapped__

    def wrapped(*args):
        """
        Wrapper for derivative of g with respect to parameter number j.
        Parameters
        ----------
        *args : iterable
            function arguments to g
        Returns
        -------
        x_bar: array-like
            derivative of g with respect to parameter number j
        """
        tans = tuple(
            [onesRobust(args[i]) if i == j else zerosRobust(args[i]) for i in range(len(args))]
        )
        _, x_bar = jvp(g, args, tans)
        return x_bar

    return wrapped