loop over variable length items

[rdemaria]

I would like to implement the following loop:

for f, args in lst:
    p=f(p,*args)

where lst can be large and f could be generic functions with variable number of arguments.

Here is a partial implementation where I restrict f to take only one numerical argument.

import jax
import jax.numpy as jnp

def fun0(p, l):
    x, px, y, py = p
    return jnp.asarray([x + px * l, px, y + py * l, py])

def fun1(p, k):
    x, px, y, py = p
    return jnp.asarray([x, px - k * x, y, py + k * y])

def fswitch(p, fa):
    fun, arg = fa
    return jax.lax.switch(fun.astype(int), [fun0,fun1], p, arg), None

@jax.jit
def loop(p, lst):
    return jax.lax.scan(fswitch, p, lst)[0]

dloop = jax.jacfwd(loop)

lst = jnp.asarray([(0, 1.2), (1, 0.8), (0, 1.2), (1, -0.7)])
p0 = jnp.asarray([0.3, 0.1, 0.2, 0.4])

p1 = loop(p0, lst)
jac = dloop(p0, lst)
print(jac)

I don't quite understand how I could extend to it to add a function like

def fun1(p, k1,k2):
    x, px, y, py = p
    return jnp.asarray([x, px - k1 * x, y, py + k2 * y])

[jakevdp]

I don't think there is any way to do this directly with scan. The best workaround would likely be to wrap your one-argument functions in functions that take two arguments, one of which is ignored: then all functions would have the same signature and you could use scan without a problem.

For example, you could write it like this:

import jax
import jax.numpy as jnp

def fun0(p, l, _):
    x, px, y, py = p
    return jnp.asarray([x + px * l, px, y + py * l, py])

def fun1(p, k1, k2):
    x, px, y, py = p
    return jnp.asarray([x, px - k1 * x, y, py + k2 * y])

def fswitch(p, fa):
    fun, *args = fa
    return jax.lax.switch(fun.astype(int), [fun0,fun1], p, *args), None

@jax.jit
def loop(p, lst):
    return jax.lax.scan(fswitch, p, lst)[0]

dloop = jax.jacfwd(loop)

lst = jnp.asarray([(0, 1.2, 0.0), (1, 0.8, 1.0), (0, 1.2, 0.0), (1, -0.7, 1.0)])
p0 = jnp.asarray([0.3, 0.1, 0.2, 0.4])

p1 = loop(p0, lst)
jac = dloop(p0, lst)
print(jac)

[rdemaria]

Thanks! For the realistic use case, I need be flexible on the arguments. Does the list passed to scan need necessarily to be a 2D array?

I tried an alternative, but tracer complains on data[iarg]...

import jax
import jax.numpy as jnp

def fun0(p, args):
    l,=args
    x, px, y, py = p
    return jnp.asarray([x + px * l, px, y + py * l, py])

def fun1(p, args):
    k1,k2=args
    x, px, y, py = p
    return jnp.asarray([x, px - k1 * x, y, py + k2 * y])

def fswitch(p, fa):
    fun, iarg = fa
    return jax.lax.switch(fun.astype(int), [fun0,fun1], p, data[iarg]), None

@jax.jit
def loop(p, ilst):
    return jax.lax.scan(fswitch, p, ilst)[0]

dloop = jax.jacfwd(loop)

lst = jnp.asarray([(0, 0), (1, 1), (0, 0), (1, 2)])
data=[[1.1],[0.8,-0.8],[-0.7,0.7]]

p0 = jnp.asarray([0.3, 0.1, 0.2, 0.4])

p1 = loop(p0, lst)

[rdemaria]

It runs, but I get the same results with different input!

[jakevdp]

Interesting - what are you seeing? When I run the code on the most recent jax and jaxlib versions, I get different results in each run:

[[ 0.11999997  1.232       0.          0.        ]
 [-0.716       0.98239994  0.          0.        ]
 [ 0.          0.          0.11999997  1.232     ]
 [ 0.          0.         -0.716       0.98239994]]
[[ 0.11999997  1.232       0.          0.        ]
 [-0.884      -0.7424      0.          0.        ]
 [ 0.          0.          0.11999997  1.232     ]
 [ 0.          0.         -0.716       0.98239994]]

i.e. it does seem to be respecting the change in global variables.

[rdemaria]

Indeed, I did not notice the difference on the second row, but the problem requires more differences.

I attach an example in which I compute with numpy the correct result and compare. The second jacobian is wrong. If I skip the first jacobian then it is correct:

import jax
import jax.numpy as jnp
import numpy as np

def fun0(p, offset, data):
    l=data[offset]
    x, px, y, py = p
    return jnp.asarray([x + px * l, px, y + py * l, py])

def fun1(p, offset,data):
    k1=data[offset+0]
    k2=data[offset+1]
    x, px, y, py = p
    return jnp.asarray([x, px - k1 * x, y, py + k2 * y])

def fswitch(p, fa):
    fun, offset = fa
    return jax.lax.switch(fun, [fun0,fun1], p, offset, data), None

@jax.jit
def loop(p, lst):
    return jax.lax.scan(fswitch, p, lst)[0]

dloop = jax.jacfwd(loop)

p0 = jnp.asarray([0.3, 0.1, 0.2, 0.4])

if True:
# first trial
    lst = jnp.asarray([(0, 0), (1, 1), (0, 0), (1, 3)])
    l=1.1;k1=0.1;k2=-0.4;
    data= jnp.asarray([l,k1,k1,k2,k2])
    p1 = loop(p0, lst)
    jac = dloop(p0, lst)

    a=np.array([[1,l,0,0],[0,1,0,0],[0,0,1,l],[0,0,0,1]])
    b=np.array([[1,0,0,0],[-k1,1,0,0],[0,0,1,0],[0,0, k1,1]])
    c=np.array([[1,0,0,0],[-k2,1,0,0],[0,0,1,0],[0,0, k2,1]])
    print(c@a@b@a-jac)

# second trial
lst = jnp.asarray([(0, 0), (1, 1), (0, 0), (1, 5)])
l=1.2;k1=0.1;k2=-0.2;
data= jnp.asarray([l,k1,k1,0,0,k2,k2])

p1 = loop(p0, lst)
jac = dloop(p0, lst)

a=np.array([[1,l,0,0],[0,1,0,0],[0,0,1,l],[0,0,0,1]])
b=np.array([[1,0,0,0],[-k1,1,0,0],[0,0,1,0],[0,0, k1,1]])
c=np.array([[1,0,0,0],[-k2,1,0,0],[0,0,1,0],[0,0, k2,1]])
print(c@a@b@a-jac)

[YouJiacheng]

Why not pass data as part of carry into scan?

[rdemaria]

It works, thanks!

[rdemaria]

Here the solution:

import jax
import jax.numpy as jnp
import numpy as np

def fun0(p, offset, data):
    l=data[offset]
    x, px, y, py = p
    return jnp.asarray([x + px * l, px, y + py * l, py])

def fun1(p, offset,data):
    k1=data[offset+0]
    k2=data[offset+1]
    x, px, y, py = p
    return jnp.asarray([x, px - k1 * x, y, py + k2 * y])

def fswitch(pdata, fa):
    p,data=pdata
    fun, offset = fa
    p=jax.lax.switch(fun, [fun0,fun1], p, offset, data)
    return (p,data), None

@jax.jit
def loop(p, lst, data):
    return jax.lax.scan(fswitch, (p,data), lst)[0][0]

dloop = jax.jacfwd(loop)

p0 = jnp.asarray([0.3, 0.1, 0.2, 0.4])

#first test
lst = jnp.asarray([(0, 0), (1, 1), (0, 0), (1, 3)])
l=1.1;k1=0.1;k2=-0.4;
data= jnp.asarray([l,k1,k1,k2,k2])

jac = dloop(p0, lst, data)

a=np.array([[1,l,0,0],[0,1,0,0],[0,0,1,l],[0,0,0,1]])
b=np.array([[1,0,0,0],[-k1,1,0,0],[0,0,1,0],[0,0, k1,1]])
c=np.array([[1,0,0,0],[-k2,1,0,0],[0,0,1,0],[0,0, k2,1]])
print(c@a@b@a-jac)

# second test
lst = jnp.asarray([(0, 0), (1, 1), (0, 0), (1, 5)])
l=1.2;k1=0.1;k2=-0.5;
data= jnp.asarray([l,k1,k1,0,0,k2,k2])

jac = dloop(p0, lst,data)

a=np.array([[1,l,0,0],[0,1,0,0],[0,0,1,l],[0,0,0,1]])
b=np.array([[1,0,0,0],[-k1,1,0,0],[0,0,1,0],[0,0, k1,1]])
c=np.array([[1,0,0,0],[-k2,1,0,0],[0,0,1,0],[0,0, k2,1]])
print(c@a@b@a-jac)

[YouJiacheng]

One more thing: pass data as carry might have performance/memory occupation issue for reverse-mode auto-diff (JAX might save data for each step). (I see you use jacfwd instead of jacrev, thus this might not be a problem, just FYI.)
If so, you can use closure and partial instead:

from functools import partial
def fun0(p, offset, data):
    l=data[offset]
    x, px, y, py = p
    return jnp.asarray([x + px * l, px, y + py * l, py])

def fun1(p, offset, data):
    k1=data[offset + 0]
    k2=data[offset + 1]
    x, px, y, py = p
    return jnp.asarray([x, px - k1 * x, y, py + k2 * y])

@jax.jit
def loop(p, lst, data):
    def fswitch(p, fa):
        fun, offset = fa
        p = jax.lax.switch(fun, [partial(fun0, data=data), partial(fun1, data=data)], p, offset, data)
        return (p,data), None
    return jax.lax.scan(fswitch, p, lst)[0][0]

BTW, your brilliant solution using flattened values array + offsets array for variable length items in JAX is great! It is a classical
layout for variable length items(used by many DBMS and libraries e.g. Apache Arrow), but I never thought it can be used with scan!