jax.lax.scan using two lists

[rdilip]

With jax.lax.scan, I take an input scan(f, init, xs). What I'd like to do is have a function of the form scan(f, init, xs, ys), where f is a two argument function such that my output is equivalent to

def scan(f, init, xs, ys):
  carry = init
  zs = []
  L = len(xs)
  assert L == len(ys)
  for i in range(L):
    carry, z = f(carry, xs[i], ys[i])
    zs.append(z)
  return carry, np.stack(zs)

So instead of scanning across a single list, I scan across the same index in two lists. It feels like this should be possible and jit-able in jax, but I'm not sure how. Any ideas?

[mattjj]

Thanks for the question! The xs argument can be any pytree of arrays, and in particular it can be a pair of arrays, so you can do this:

import jax.numpy as jnp
from jax import lax

xs = jnp.array([1, 2, 3])
ys = jnp.array([3, 1, 4])

def scanned_fun(_, pair):
  x, y = pair
  return None, x + y

_, zs = lax.scan(scanned_fun, None, (xs, ys))
print(zs)  # [4, 3, 7]

We're ignoring the carry here and in effect just mapping the scalar + function over two vectors.

WDYT?

[rdilip]

Thanks for the quick response! That all makes sense. As a follow up, is there a way to do this with lists of different sizes? More specifically, here's what I'd like to do.

xs is a list of tensors each with shape (for example) (2, 5, 5)
ys is a list of tensors each with shape (2,2,2,2)

I would like to do the following for each index i through L-1

def f(xs, ys, i):
    x0, x1 = xs[i], xs[i+1]
    y0 = ys[i]
    combined_x = np.tensordot(x0, x1, [2,1]) # has shape (2, 5, 2, 5)
    y_on_x = np.tensordot(y0, combined_x, [[0,1],[0,2]]) # has shape (2, 2, 5, 5)
    y_on_x = y_on_x.transpose((0,2,1,3)).reshape((2*5, 2*5))
    Q, R = np.linalg.qr(y_on_x)
    Q = Q.reshape((2, 5, -1))
    R = R.reshape((-1, 2, 5)).transpose((1,0,2))
    xs[i] = Q
    xs[i+1] = R

So in summary, I want to take pairs of one list and contract them with individual tensors from the second, and iterate down in this way. Is this possible?

[mattjj]

Scanning over arrays of different shapes is fine, so long as they have the same leading axis size. That is, the array leaves on the input pytree all have to have the same leading axis size. (Did I misunderstand your question?)

There's another thing to handle in your example, which is that each iteration of the loop actually writes to xs[i+1]. You might want to lift that value into the carry. That is, the carry could include the R factor from the previous iteration. I haven't worked out the details myself though.

[rdilip]

Sorry for the delay -- just noted this response. In this specific example, though, it seems impossible because the two arrays would have different leading axis shapes. I'm asking about a case where I have one array with L elements and one array with L-1 elements. I contract each consecutive pair of the first array with each element of the second array (but do so sequentially). However, this means the leading axis of the two arrays will differ by 1. Is there a way to address this problem?

Re the second point you raise, that makes sense -- I think I see how to do it if I can resolve the first issue.