Apply function (jnp.var) on all sub-segments of a 1D array (TracerIntegerConversionError)

[deepcharles]

Hi,

Thanks for the amazing work. I find it very useful.

I have a problem where I need to compute the variance of all subsegments of a 1D array, i.e. signal[start:end].var() for all start, end such that 0 <= start < end <= n_samples when the shape of signal is (n_samples,).
This seems like an easily parallelizable code but I cannot make it work somehow. Please find some code below.

Can you think of a way to make this code work (the second attempt mostly)?

import jax.numpy as jnp
import jax.random as jrd
from jax import lax
from jax.typing import ArrayLike as JArrayLike

# generate dummy signal
n_samples = 100
seed = 123
key = jrd.PRNGKey(seed)
signal = jrd.normal(key=key, shape=(n_samples,))


# Compute the variance of all subsegments

# First try, naive numpy translation.
# does not work because of assignment
res = jnp.zeros((n_samples + 1, n_samples + 1))
for start in range(n_samples):
    for end in range(start + 1, n_samples + 1):
        res[start, end] = signal[start:end].var()  # assignment

# Second try, not sure why this does not work.
# does not work, raise a TracerIntegerConversionError:
# TracerIntegerConversionError: The __index__() method was called on the JAX Tracer object Tracedwith with
#   val = Array([1, 2, 3, ..., 1, 2, 1], dtype=int32)
#   batch_dim = 0
# See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerIntegerConversionError

def get_single_var(start: jnp.int32, end: jnp.int32) -> jnp.float32:
    sub = lax.dynamic_slice_in_dim(
        signal, start_index=start, slice_size=end - start
    )
    return jnp.var(sub)

row_indexes, col_indexes = jnp.triu_indices(n_samples + 1, k=1)
res = jnp.vectorize(get_single_var)(row_indexes, col_indexes)

Cheers,

Charles

[deepcharles]

After looking at #2323, I suspect that this task might not be efficiently solved with jax. But I'd like your opinion on that.

[jakevdp]

Hi - thanks for the question. Working with algorithms that involve dynamically-shaped arrays sometimes requires some thought, but is often still possible to do efficiently.

Your first attempt is running into JAX's immutable arrays, but you can use the functional in-place update syntax intead (see Sharp Bits: In Place Updates):

import jax.numpy as jnp
from jax import random

n_samples = 4
seed = 123
key = random.PRNGKey(seed)
signal = random.normal(key=key, shape=(n_samples,))


def f_loop(signal):
  n_samples = len(signal)
  res = jnp.zeros((n_samples + 1, n_samples + 1))
  for start in range(n_samples):
    for end in range(start + 1, n_samples + 1):
      res = res.at[start, end].set(signal[start: end].var())
  return res

f_loop(signal)

Array([[0.        , 0.        , 1.5675846 , 3.1089766 , 2.3345966 ],
       [0.        , 0.        , 0.        , 4.621629  , 3.0940673 ],
       [0.        , 0.        , 0.        , 0.        , 0.91021997],
       [0.        , 0.        , 0.        , 0.        , 0.        ],
       [0.        , 0.        , 0.        , 0.        , 0.        ]],      dtype=float32)

If you'd like to do the same thing more efficiently using vmap, one way is to write a simple dynamic_var utility that will dynamically compute the variance of a slice of an array. Here's one possible approach, generalizing the strategy mentioned in Sharp Bits: Dynamic Shapes:

def dynamic_var(x, start, end):
  """Compute x[start:end].var() with dynamic start and end"""
  i = jnp.arange(len(x))
  x = jnp.where((i >= start) & (i < end), x, 0)
  return jnp.sum(x ** 2) / (end - start) - jnp.sum(x) ** 2 / (end - start) ** 2

def f_vmap(signal):
  n_samples = len(signal)
  i, j = jnp.triu_indices(n_samples + 1, k=2)
  res = jnp.zeros((n_samples + 1, n_samples + 1))
  return res.at[i, j].set(jax.vmap(dynamic_var, (None, 0, 0))(signal, i, j))

f_vmap(signal)

Array([[0.       , 0.       , 1.5675848, 3.1089764, 2.3345964],
       [0.       , 0.       , 0.       , 4.621629 , 3.094067 ],
       [0.       , 0.       , 0.       , 0.       , 0.9102199],
       [0.       , 0.       , 0.       , 0.       , 0.       ],
       [0.       , 0.       , 0.       , 0.       , 0.       ]],      dtype=float32)

This second approach should be much faster than the simpler approach based on nested for-loops, particularly as n_samples grows large.

[deepcharles]

Thanks, that's perfect.