Add optax.tree_utils.tree_batch_shape.

[carlosgmartin]

This commit adds a batch_shape argument to optax.tree_utils.tree_random_like.

This is useful when sampling multiple perturbations, such as for evolution strategies.

[carlosgmartin]

@rdyro How does this look? Happy to make any needed changes.

[rdyro]

Hey, sorry for the delayed response. Adding batch is definitely useful in practice, but I'm concerned about API uniformity with the other tree_like methods:

tree_full_like
tree_ones_like
tree_random_like
tree_split_key_like
tree_zeros_like

Can we go with a workaround for the time being (I realize it's not as concise 😞 )

a = {"a": jnp.ones(10), "c": jnp.zeros(4), "d": jnp.array(1.0)}
bs = (10,)
batch_tree = jax.tree.map(lambda x: jax.ShapeDtypeStruct(bs + x.shape, x.dtype), a)
optax.tree_utils.tree_random_like(random.key(0), batch_tree)

[carlosgmartin]

How about adding a helper function like the following?

def tree_batch(a, shape, dtype=None, sharding=None):
  return jax.tree.map(lambda a: jax.ShapeDtypeStruct(
    shape + a.shape,
    dtype=dtype or a.dtype,
    sharding=sharding or a.sharding,
  ), a)

Then we can use

from jax import random
from optax import tree_utils as otu
otu.tree_random_like(random.key(0), otu.tree_batch(tree, shape))

[rdyro]

That sounds like a great idea! Would you be willing to contribute this addition?

Can you call it something a little more verbose than tree_batch maybe tree_batch_shape, but I'm open to better ideas?

[carlosgmartin]

@rdyro Done. I thought it might be more helpful for other potential uses to keep the actual values instead of returning ShapeDtypeStructs, so I've switched to using broadcast_to instead. (This uses zero-strides, so it doesn't need additional memory.)

[carlosgmartin]

@rdyro Is the above ok?

[rdyro]

@rdyro Is the above ok?

Using broadcast_to makes this function more general, nice!