Make jax._src.api._std_basis part of the Main tree_utils API

[joeryjoery]

Hey, I've now repeatedly come across situations in which a jnp.eye function generalized to pytrees would be quite helpful. Though, it surprised me that in the backend there is a _std_basis function for pytrees that is hidden from the main API. Why is this the case?

I'd say that it would be helpful to add this as part of the main API, maybe as part of the jax.tree_util API?

An example use case I had was for evaluating each intermediate Jacobian within a function composition; or in other terms, the adjoint VJPs at each discrete time-step t fed through the standard basis.

from functools import reduce

import jax
import jax.numpy as jnp

import chex

f = lambda *a: jnp.tile(jnp.squeeze(jnp.asarray(a)), 2) ** 2
g = lambda a: a[::2] / jnp.sqrt(a[1::2])

funs = [f, f, f, g]
F = reduce(lambda a, b: (lambda *z: b(a(*z))), funs)

x = (12., 3.1, 0.8)
print(F(*x)) 
# [4.6558955e+06 8.0908667e-06 2.0822533e+04 4.6558955e+06 8.0908667e-06
# 2.0822533e+04 4.6558955e+06 8.0908667e-06 2.0822533e+04 4.6558955e+06
# 8.0908667e-06 2.0822533e+04]

z = x
orbit = [jax.device_put(z)]
for f in funs:
    z = f(z)
    orbit.append(z)
    
print(jax.tree_map(jnp.shape, orbit))
# [((), (), ()), (6,), (12,), (24,), (12,)]

def orbit_adjoint(primals, cotangents):
    orbit_pullback = list()
    for t in reversed(range(len(funs))):
        pullback = jax.vjp(funs[t], primals[t])[1]
        cotangents = pullback(cotangents)[0]

        orbit_pullback.append(cotangents)

    return (orbit_pullback[::-1], )

# ! std-basis needed
adjoints, = jax.vmap(orbit_adjoint, in_axes=(None, -1))(orbit, jax._src.api._std_basis(orbit[-1]))  
print(jax.tree_map(jnp.shape, adjoints))
# [((12,), (12,), (12,)), (12, 6), (12, 12), (12, 24)]

# The first index of `adjoints` contains `dF/dx`.
chex.assert_trees_all_close(adjoints[0], jax.jacrev(F, argnums=(0, 1, 2))(*x))

Don't pay too much attention to the function F itself, it's just an example. But this would be useful for Neural Networks in e.g., layer-wise (block-wise) factorizations of the Hessian/ Fisher.

This is one example where the std-basis is helpful. It can also be helpful for eliminating off-diagonal elements in any ArrayTree, or for defining complex matrix vector products with pytrees.

What do you all think?

[patrick-kidger]

Worth noting that _std_basis isn't quite a generalisation of jnp.eye in the way that might be expected. (Or at least that I'd expect.)

jax._src.api._std_basis([jax.numpy.zeros(3), jax.numpy.zeros(3)])

produces

[DeviceArray([[1., 0., 0.],
             [0., 1., 0.],
             [0., 0., 1.],
             [0., 0., 0.],
             [0., 0., 0.],
             [0., 0., 0.]], dtype=float32), DeviceArray([[0., 0., 0.],
             [0., 0., 0.],
             [0., 0., 0.],
             [1., 0., 0.],
             [0., 1., 0.],
             [0., 0., 1.]], dtype=float32)]

which is just a length-2 list of arrays. It is not a length-2-list-of-length-2-lists of arrays. That is to say the "PyTree axes" don't get duplicated out the way the "array axes" do.

It's very easy to write a std_basis function that does this alternate (more straightforward?) thing, though:

import jax
import jax.numpy as jnp

def std_basis(pytree):
    leaves, structure = jax.tree_flatten(pytree)
    eye_structure = structure.compose(structure)
    eye_leaves = []
    for i1, l1 in enumerate(leaves):
        for i2, l2 in enumerate(leaves):
            if i1 == i2:
                eye_leaves.append(jnp.eye(jnp.size(l1)).reshape(jnp.shape(l1) + jnp.shape(l2)))
            else:
                eye_leaves.append(jnp.zeros(l1.shape + l2.shape))
    return jax.tree_unflatten(eye_structure, eye_leaves)

and of course any time you need this "matrix PyTree" to interact with a "vector PyTree" you can do so by just jax.tree_map'ing them together.

[joeryjoery]

Thanks for the comment! Yes, I think your approach makes sense if I would want to do vector-matrix like arithmetic, so this would be ideal for use with tree-map. The jax private API is quite useful for doing things with vmap as @YouJiacheng mentioned.

[joeryjoery]

Maybe that's why I was also wondering to why the private _std_basis is actually private. Is it because of such a confusion?

[YouJiacheng]

As @patrick-kidger has shown, there are two kinds of "std_basis":

1. has the same tree structure as input-tree while leaf prepended with a leading axis of size n := sum(leaf.size for leaf in jax.tree_util.tree_flatten(tree)[0]), which is useful for vmap + jvp or vjp
2. has nested tree structure as patrick-kidger described.

[YouJiacheng]

BTW, you can implement jax._src.api._std_basis using only public API in 2 lines:
Considering the ambiguity of std_basis, adding it to autodiff cookbook or some other advanced tutorials seems better.

import jax
import jax.numpy as jnp
from jax.flatten_util import ravel_pytree

def std_basis_for_vmap(pytree):
    flat, unravel_pytree = ravel_pytree(pytree)
    return jax.vmap(unravel_pytree)(jnp.eye(jnp.size(flat), dtype=jnp.dtype(flat)))

print(std_basis_for_vmap([jnp.zeros(3), jnp.zeros(3)]))

[joeryjoery]

Thanks! That's indeed a good illustration.