jit-able function that can loop a matrix over its diagonal lines?

[riven314]

I have a function that constructs a matrix and then loop over its diagonal axis to update the entries. (e.g. loop from top-left diagonal line till bottom-right diagonal line).

In numpy, it is something like this:

r, c = np.array(D.shape)-1
for a in range(1, r+c):
    # I, J are the row and col indexes of the diagonal line to be accessed
    I = np.arange(max(0, a-c), min(r, a))
    J = I[::-1] + a - min(r, a) - max(0, a-c)
    # two np.minimum because np.minimum takes only two args.
    D[I+1, J+1] += np.minimum(np.minimum(D[I, J], D[I, J+1]), D[I+1, J])

However, I have hard time converting the logic to jax jit-able code. Main bottleneck I found is that the length of the indices (i.e. I, J) that I use for indexing the diagonal line depends on the jax operations (i.e. min, max). So it seems the logic above cant be directly translated to jit-able code.

Do you have any advice to rephrase the logic so that it is jit-able? Thanks!

[C-J-Cundy]

Hi, here is a quick solution I whipped up, which I think works. To be honest I'm not quite sure what your operation is doing, and I suspect there might be a way to re-write it with masks so that it doesn't take so long to compile.

As for general tips for re-writing, I find that it helps to have a clear idea of what the jit compiler has access to and what it doesn't. In this case, we are allowed to use I and J inside a jit function, as they are fixed given the shape of the input (and the jit compiler knows the shape of the input when compiling). So we can use quantities like np.shape(X), we can write loops where the numbers of iterations depend on the shape or quantities derived from the shape, etc.

An annoying detail is that the compiler is typically super-linear in the length of python loops, which are statically unrolled when compiled. So the advice is generally to avoid long python loops if possible. Here there is likely some way to avoid the loops via a scan, but this will require additional rephrasing. If you're only dealing with a fixed small r,c, and not recompiling often, hopefully you can deal with the compile time!

import jax.numpy as jnp
import jax.random as rnd
import numpy as np
import jax
import time
np.random.seed(0)

def do_op_numpy(D):
  assert type(D) is np.ndarray
  r, c = np.array(D.shape)-1
  for a in range(1, r+c):
    # I, J are the row and col indexes of the diagonal line to be accessed
    I = np.arange(max(0, a-c), min(r, a))
    J = I[::-1] + a - min(r, a) - max(0, a-c)
    # two np.minimum because np.minimum takes only two args.
    D[I+1, J+1] += np.minimum(np.minimum(D[I, J], D[I, J+1]), D[I+1, J])
  return D

@jax.jit
def do_op_jax(D):
  r, c = np.array(D.shape)-1
  for a in range(1, r+c):
      # I, J are the row and col indexes of the diagonal line to be accessed
      I = np.arange(max(0, a-c), min(r, a), dtype=jnp.int32)
      J = I[::-1] + a - min(r, a) - max(0, a-c)
      summand = jnp.min(jnp.stack([D[I, J], D[I, J+1], D[I+1, J]]), axis=0)
      D = jax.ops.index_add(D, jax.ops.index[I+1, J+1], summand)
  return D

# check output is the same for original and new implementation
n_trials = 3
for n in range(n_trials):
  r, c = np.random.randint(low=1, high=100, size=2)
  D = np.random.normal(size=(r,c))
  D_jax = jnp.array(D)
  np_output = do_op_numpy(D)
  # Takes a long time to compile!
  jax_output = do_op_jax(D_jax)
  assert np.allclose(np_output, jax_output, rtol=1e-4, atol=1e-6)
  #assert np.allclose(np_output, jax_output)

# check speed
r, c = 400, 400
D = np.random.normal(size=(r,c))
t0 = time.time()
np_output = do_op_numpy(D)
t1 = time.time()
numpy_time = t1 - t0

# Compile
t0 = time.time()
jax_output = do_op_jax(jnp.array(D))
t1 = time.time()
jax_output = do_op_jax(jnp.array(D))
t2 = time.time()

print(f"For {(r,c)}-dimensional input, took {numpy_time}s for numpy, {t2 - t1}s for JAX, compiled in {t1 - t0}s on GPU colab")
#For (400, 400)-dimensional input, took 0.024621248245239258s for numpy, 0.030493497848510742s for JAX, compiled in 77.43533253669739s on GPU colab

[riven314]

Thanks for the reply! With reference to your code, I changed my code like this:

@jax.jit
def dynamic_time_warping(s1, s2):
    r = s1.shape[0]
    c = s2.shape[0]
    D = jnp.full((r+1, c+1), jnp.inf)
    D = jax.ops.index_update(D, jax.ops.index[0, 0], 0.)
    for a in range(1, r+c):
        I = jnp.arange(max(0, a-c), min(r, a))
        J = I[::-1] + a - min(r, a) - max(0, a-c)
        min_val = jnp.minimum(jnp.minimum(D[I, J], D[I, J+1]), D[I+1, J])
        cost = min_val + jnp.square(s1[I]-s2[J])
        D = jax.ops.index_update(D, jax.ops.index[I+1, J+1], cost)
    return D[-1, -1]

I just learned it is valid for a jit-able function to have jnp.arange(max(0, a-c), min(r, a)). I confused it with the case of jnp.maximum, jnp.minimum here coz these jax operations are not allowed in jnp.arange.

However, I found that the jit compiling time is super slow when s1, s2 have a long length (e.g. (600,)). I tried to wrap the whole loop by loops.Scope() but the issue still persist. I am trying to replace the python for loop by jax.lax.fori_loop or jax.lax.scan. Do you think it is feasible in this case? cause the loop here is mainly doing index_update.

[riven314]

The difficulty I had when writing it to fori_loop is that it will convert a into a traced object, and hence invalidate the operation jnp.arange(max(0, a-c), min(r, a)). If I convert it to jnp.arange(jnp.maximum(0, a-c), jnp.minimum(r, a)), then it cant be jit-able

[C-J-Cundy]

The key issue here is that you can't use inputs/outputs with variable shape for fori_loop or scan. What I would probably do is to use a scan, where we carry through the current and previous antidiagonal, both as length-d vectors which are masked (i.e., padded smaller vectors).

I'm not sure if this is what you're looking for exactly, but it appears there is already an implementation of DTW for jax available here. It seems that they indeed use this approach.