Question about achieving certain behaviour with lax.scan

[adam-hartshorne]

I can't figure out how to achieve the following using lax.scan, where I commented out code is the behaviour I desire (but obviously I am aware you can't use if / else inside the scan). Any help much appreciated.

yn = [y1]
init_val = (y0, yn)

def scan_fn(carry, y):
    y0, y1 = carry

    # if len(y) == 0:
    #     yn = [y1]
    # else:
    #     yn = y

    y_next, yn_out = jet.jet(self.dynamics, (state, ), ((y0, *yn),),)
    return carry, yn_out

_, yn = jax.lax.scan(scan_fn, init=init_val, length=order)

[jakevdp]

Can you say more about what your inputs are and what your expected output is? It's difficult to tell what you're attempting to do without more information. Consider editing your question to add a minimal reproducible example. Thanks!

[adam-hartshorne]

Here is basically what it is doing. yn is just a list that each time around the loop increases in length by 1.

order = 4
y0 = jnp.zeros((3,2))
y1 =[jnp.zeros((3,2)),]

carry = (y0, y1)

for _ in range(order-1):
    y0, yn = carry
    # function call y0, *yn = jet.jet(self.dynamics, (state,), ((y0, *yn,),)) 
    # replaced with a placeholder that results in the same dimensions.
    yn.append(jnp.zeros((3,2)))
    carry = (y0, yn)

y0, yn = carry
output = jnp.vstack([y0[jnp.newaxis], yn])

The problem is we can't do the following because yn is a list increasing in size, and because the updated yn is required for the real function call at each pass around the loop, storing the updated yn in the second return argument (currently set to None) is of no use.

order = 4
y0 = jnp.zeros((3,2))
y1 =[jnp.zeros((3,2)),]

def update_state(carry, unused_input):
    y0, yn = carry
    yn.append(jnp.zeros((3,2)))
    return (y0, yn), None

(y0, *yn), _ = jax.lax.scan(update_state, ((y0, y1)), xs=None, length=order-1)

output = jnp.vstack([y0[jnp.newaxis], yn])

[jakevdp]

I see - thanks for making the intent more clear. Since your intent is to append a new value each iteration, maybe something like this is what you're looking for?

import jax.numpy as jnp
import jax

def update(y0, yn):
  yi =  jnp.zeros((3, 2))
  return y0, yi

y0, yn = jax.lax.scan(update, y0, xs=None, length=order)
output = jnp.vstack([y0[jnp.newaxis], yn])
print(output.shape)
# (5, 3, 2)

[adam-hartshorne]

That won't work, because I need to pass both initial values for y0 and yn to begin with, as the following function is called in my real code,

y0, *yn = jet.jet(self.dynamics, (state,), ((y0, *yn,),))

[jakevdp]

Oh, in that case then I think you're out of luck: there's no way to use scan for computations which have different array sizes each iteration. You may be ble to make progress by using a static-sized yn padded with zeros, but that would involve rewriting your self.dynamics function to handle that kind of padded input: I'm not sure how feasible that would be.

[pnkraemer]

Sorry to jump in from the side like this, but what @adam-hartshorne is trying to achieve here (the code snippets seem to represent recursively computing ODE solution coefficients with jet) is indeed possible with scan and a little bit of padding :)

@adam-hartshorne and I discussed this in a different issue (in the repo linked below), so I think he is already aware of this, but for anyone else encountering the same problem: I implemented a solution in the taylor_mode_fn here which is part of probdiffeq.