New pytree via sampling leaves of original pytree

[pharringtonp19]

Aim: Given a pytree, I would like to construct a new pytree by sampling from the leaves of the given pytree

I was thinking that something like the following might work. I'm not sure though how to derive a new PyTreeDef from the original one.

def sample_tree(key, tree, n):
  idx = jax.random.choice(key, len(tree), shape=(n,), replace=False).astype(jnp.int32)
  new_flat = [jax.tree_util.tree_leaves(tree)[i] for i in idx]
  new_tree = ...?...
  return jax.tree_util.tree_unflatten(new_tree, new_flat)

[jakevdp]

I think we need more information here: what is the expected structure of your output? e.g. say tree is a a nested dict of dicts, and you sample n of the leaf values. What should the output look like?

[pharringtonp19]

@jakevdp Apologies, I did have a specific structure in mind. I've written down a toy example of my problem (I've borrowed some of this from jax stateful tutorial).

Context: My data set consists of clusters of observations where the number of observations in each cluster can vary.

Specific Aim: In each training step, I would like train using only a subset of the clusters.

Apporaches: My current approach does this by applying a mask in the loss function. I would like to compare this approach to (sampling a subset of the clusters and then tree-mapping across only those sampled clusters). My understanding is that if I compile the loss function, then these approaches should be similar? I represent my dataset via a dictionary of jnp.arrays, but so far I have been unable to successfully implement this sampling approach

Masked Loss Approach

import jax 
from jax.tree_util import tree_flatten, tree_unflatten 
import jax.numpy as jnp 
from functools import partial 
from typing import NamedTuple
from collections import OrderedDict

def sample(key, n):
  return jax.random.normal(key, shape=(n,2))

class Params(NamedTuple):
  weight: jnp.ndarray
  bias: jnp.ndarray

def init(rng) -> Params:
  """Returns the initial model params."""
  weights_key, bias_key = jax.random.split(rng)
  weight = jax.random.normal(weights_key, ())
  bias = jax.random.normal(bias_key, ())
  return Params(weight, bias)

def per_cluster_loss(params: Params, data: jnp.ndarray) -> jnp.ndarray:
  x, y = data[:,0].reshape(-1,1), data[:,1].reshape(-1,1)
  """Computes the least squares error of the model's predictions on x against y."""
  pred = params.weight * x + params.bias
  return jnp.mean((pred - y) ** 2)

def loss_fn(params: Params, data: dict, key) -> jnp.ndarray:
  cluster_losses = jax.tree_util.tree_map(partial(per_cluster_loss, params), data)
  cluster_mask_idx = jax.random.permutation(key, jnp.repeat(jnp.array([1., 0.]), len(data)//2))
  cluster_mask = {i: cluster_mask_idx[i] for i in data.keys()}
  masked_cluster_losses = jax.tree_util.tree_map(lambda p, q: p*q, cluster_losses, cluster_mask)
  loss = (1/(len(data)/2))*jax.tree_util.tree_reduce(lambda a, b: a +b, masked_cluster_losses)
  return loss

params = init(jax.random.PRNGKey(0))

data = {}
key = jax.random.PRNGKey(0)
for i in range(10):
  key, subkey1, subkey2 = jax.random.split(key, 3)
  n = jax.random.choice(subkey1, jnp.array([3,5,7]))
  data[i] = sample(subkey2, n)

print(loss_fn(params, data, jax.random.PRNGKey(0)))

Failing Sampling Approach

def loss_fn2(params: Params, data: dict, key) -> jnp.ndarray:
  idx = jax.random.choice(key, jnp.arange(len(data)), shape=(5,), replace=False)
  mini_batch = {i: data[i] for i in idx}
  cluster_losses = jax.tree_util.tree_map(partial(per_cluster_loss, params), mini_batch)
  loss = (1/(len(data)/2))*jax.tree_util.tree_reduce(lambda a, b: a +b, cluster_losses)
  return loss

     36 def loss_fn2(params: Params, data: dict, key) -> jnp.ndarray:
     37   idx = jax.random.choice(key, jnp.arange(len(data)), shape=(5,), replace=False)
---> 38   mini_batch = {i: data[i] for i in idx}
     39   cluster_losses = jax.tree_util.tree_map(partial(per_cluster_loss, params), mini_batch)
     40   loss = (1/(len(data)/2))*jax.tree_util.tree_reduce(lambda a, b: a +b, cluster_losses)

TypeError: unhashable type: 'DeviceArray'

[jakevdp]

I see, thanks. You're not going to be able to reference dict items with traced values. Fundamentally, the issue is that at trace-time JAX needs to know the size of data[i], but if i is not determined until runtime, there's no way to determine this!

The best approach would probably be to stick those entries into a single 2D array rather than a dict. This would be easy if all arrays were the same size, but with different sizes you may need to use a padded array instead, and rework your code to take into account that padding.

There is ongoing work on dynamic shapes that may make this kind of thing easier in the long run, but right now it might require some manual tinkering. Does that make sense?

[pharringtonp19]

@jakevdp Thanks for your help again

JAX needs to know the size of data[i], but if i is not determined until runtime, there's no way to determine this!

Makes perfect sense -- Thanks