How can I do arithmetic with gradients that I get from Stax?

[geajack]

Let's say I make a neural network in Stax:

initializer, model = stax.serial(
    stax.Dense(1000),
    stax.Relu,
    stax.Dense(1),
    stax.Relu
)

and then later I'm training this thing, so I have some code like this:

grad_L = jit(grad(lambda w, x, y: L(model(w, x), y)))
gradient = grad_L(w, x, y)

Now I want to get the norm of the gradient. I might try:

np.sqrt(np.sum(gradient**2))

Or maybe I want to do arithmetic with this thing, such as multiplying it by a scalar, maybe because I'm implementing my own optimization algorithm:

0.1 * gradient

Of course, neither of these snippets will work, because gradient isn't a numpy array or anything similar: it's a Python tuple of Python tuples of DeviceArrays. So to do either of these operations, I have to write a nested for loop, which possibly might even break if I change the network architecture.

Is there any way of doing arithmetic on the result of a grad function? I'm hoping for an answer in one of three forms:

1. A set of helper functions in JAX for doing arithmetic on nested Python containers, like scaled_gradient = jnp.multiply(0.1, gradient), or something.
2. A way to efficiently convert from and to this nested format, something like vector = jax.vectorize(gradient), new_gradient = jax.unvectorize_like(vector, gradient).
3. A way to set up my network with Stax so that its weights are in a vector format to begin with, and grad returns nice arithmetically compatible objects in the first place.

[mattjj]

Thanks for the question!

The easiest thing to do use use tree utilities, basically mapping operations:

# corresponding to #1 of your example, this is the most canonical
from jax.tree_util import tree_map
scaled_gradient = tree_map(lambda x: 0.1 * x, gradient)

# corresponding to #2
from jax.flatten_util import ravel_pytree
vector, unflatten = ravel_pytree(gradient)
scaled_gradient = unflatten(0.1 * vector)

There's a way to do the third option too. But the first option seems the most canonical. You might even want to just work with tree_util.tree_flatten and tree_util.tree_flatten just to get flat lists and map over them, since that's what tree_map does under the hood!

That said, keep in mind that Stax is just an example library for inspiration. It's not actively developed, though people sometimes fork it to write their own version for their own needs.

If you want better libraries, take a look at these.

[geajack]

Great, thanks. How would you recommend computing the norm of the gradient? It seems to me that that problem doesn't quite fit the shape of your example code. Right now I'm using this:

squared_components, _ = tree_flatten(tree_map(lambda u: jnp.sum(u**2), gradient))
squared_norm = jnp.sum(jnp.array(squared_components))
gradient_norm = jnp.sqrt(squared_norm)

Is that idiomatic? Any tips to improve speed?