Nested Pytree Indexing

[joeryjoery]

Say I have a Pytree a, which is an outer product of b with b.

How could I efficiently extract elements from a given the pytree definition of b?
Because the structure of a is defined by the structure of b for every element in b.

An example use-case is to extract the diagonal of the Hessian contained in an awkard nested pytree without flattening.

[YouJiacheng]

What do you mean by "extract"?

[joeryjoery]

With extract, I generally meant look-up or reading. Just accessing the item for its value.

Funnily enough, while reviewing my question and debugging for longer I kind of figured out what I wanted to do.

What I wanted to do can be achieved through the Pytree leaves (pseudo-Jax code):

a = outer(b, b)
flat = jax.tree_leaves(a)
b_like_flat = [flat[c] for c in my_idxs]
b_like = jax.tree_unflatten(jax.tree_structure(b), b_like_flat)

I needed this for e.g., the Hessian diagonal,

my_idxs = (len(b) + 1) * jnp.arange(len(b))

It's still inefficient storage-wise though, because this method still requires materializing a.

[YouJiacheng]

Hmmm, I think that you code snippet will result in b_like == b?

[joeryjoery]

If a is the outer product of b and b and my_idxs are the diagonal indices.

Then b_like == square(b), where square is elementwise.

I.e., taking the sum of b_like then gives the squared l2 norm of b since the trace of an outer product is equal to the inner product.

[YouJiacheng]

Oh, I don't notice that outer not only affects structure, but also implies multiplication.
For this use case, maybe you can just use for-loop to access flat b with 2-d indices.
Since indexing applies to pytree, any indexing method should result in similar code after jit. Then simple for-loop just win!

idx1 = range(len(b))
idx2 = range(len(b))
b_flat = jax.tree_leaves(b)
b_like_flat = [b_flat[i] * b_flat[j] for i, j in zip(idx1, idx2)]
b_like = jax.tree_unflatten(jax.tree_structure(b), b_like_flat)