Numpy loops to Jax: conversion help

[alonfnt]

I'm having trouble converting a numpy code to Jax and I'm hoping the community can help me out.

Here's the numpy code:

import numpy as np

t, n = 10, 5
x = np.random.normal(size=(t, n))
r = np.random.uniform(size=(t, t))
z = np.zeros((t, t, n))
for t1 in range(t - 1):
    y = x[t1]
    z[t1, t1 + 1] = np.copy(y)
    for t2 in range(t1 + 1, t - 1):
        y = y + x[t2] * (r[t2, t1:t2] @ z[t1, t1:t2])
        z[t1, t2 + 1] = np.copy(y)
result = z.sum(axis=-2)

I've tried to use jax.vmap on t1 since there are no dependencies between loops, but I'm struggling with the second loop. The expanding window component and the previous calculation being included in the z array make it difficult to convert to Jax:

@jax.vmap
def f(xi, t1): # This code fails due to t1 going into range
    z = jnp.zeros(xi.shape)
    for t2 in range(t1+1, t-1):
        xi = xi + x[t2] * (r[t2, t1:t2] @ z[t1:t2])
        z[t1, t2 + 1] = xi
    return z

result = f(x, jnp.arange(t)).sum(axis=-2)

Any help would be greatly appreciated!

[jakevdp]

This is not an easy function to express in terms of JAX transformations or higher-order primitives. You can't use vmap, because your function is not a purely batch-wise function (the result at z[t1] depends on all x[t2] where t2 > t1)

Further, you won't be able to directly express this within lax.fori_loop or similar because each iteration constructs dynamically-shaped intermediate arrays (if you're iterating over t1 and t2 using fori_loop, the size of z[t1:t2] will be dynamic).

That said, you could use this within JAX nearly as-written if you modify it to use JAX's functional array updates:

import jax

@jax.jit
def f(x, r, z):
  t, n = x.shape
  assert r.shape == (t, t)
  assert z.shape == (t, t, n)
  for t1 in range(t - 1):
    y = x[t1]
    z = z.at[t1, t1 + 1].set(y)
    for t2 in range(t1 + 1, t - 1):
        y = y + x[t2] * (r[t2, t1:t2] @ z[t1, t1:t2])
        z = z.at[t1, t2 + 1].set(y)
  return z.sum(axis=-2)

result_jax = f(x, r, z)
print(result_jax)

As long as t is relatively small, this should be fine, but as t grows, you can expect compilation time to scale roughly as $\sim t ^ 4$ (quadratic in the number of statements, which is $\mathcal{O}[t^2]$) so this could end up being too slow.

If that compilation time for the flattened loops is problematic, you best bet is probably to use a nested fori_loop, but figure out how to express the dynamic matrix product in terms of static shapes, which should be doable with a bit of thought.

[alonfnt]

I see... using jit with the loops takes too much time to compaile, so that's not an option. I'll have to think more about how to deal with the shapes -- maybe taking advantage of the zeros can make them less dynamic.

Thanks for the insight!

[soraros]

I rewrote the program to one with a simplified loop structure. Given another rainy afternoon, I may finish the rewrite and come up with a fully vectorized version. Yet here is it now.

def f_np(x: npt.NDArray, r: npt.NDArray) -> npt.NDArray:
  t, n = x.shape
  assert r.shape == (t, t)
  z = np.zeros((t, t, n))
  I, J = np.triu_indices(t, 1)
  z[I, J] = x[I]
  I, J = np.triu_indices(t - 1, 1)
  for i, j in zip(I, J):
    z[i, j + 1:] += x[j] * (r[j, :j] @ z[i, :j])
  return z.sum(axis=-2)