Trivial dynamic shapes

[awav]

Hi all,

Is there a way to have dynamic shapes, slices and other objects with dynamic properties in Jax for trivial cases. I would appreciate it if anyone could guide to the right solution. Thanks!

The minimal (non)working example is below. It appears to me that this case is trivial, but very surprising that it doesn't work in jax. The index i has a maximum value and the increment is linear, there is a "clear" computational dependency on the index i.

def cond_fn(acc):
    i, _ = acc
    return i < 5

def body_fn(acc: tuple[int, jax.typing.ArrayLike]):
    i, var = acc
    update = jnp.ones([2, i + 1]) * i # << Breaks here, and even jnp.repeat is required to be static...
    var_new = lax.dynamic_update_slice(var, update, [i, 0])
    return i+1, var_new
    
lax.while_loop(cond_fn, body_fn, (0, arr))

PS: it seems that even lax.dynamic_update_slice doesn't support dynamic indices. In my case it is [i, 0].
Relevant: #14634

[jakevdp]

No, dynamic shapes are not yet supported, even trivial ones that you mention (in practice, this is not so trivial to implement!) You linked to #14634, I think that's probably the place to follow for updates.

Regarding lax.dynamic_update_slice, it does support dynamic indices, but not dynamic array shapes. This is in contrast to lax.slice, which does not support dynamic indices at all, e.g.

import jax

@jax.jit
def get_slice(x, i):  # note: i is dynamic
  # return jax.lax.slice(x, [i], [i + 3])  # this would fail
  return jax.lax.dynamic_slice(x, [i], [3])  # this succeeds

get_slice(jax.numpy.arange(5), 2)

[awav]

Thank you for the swift response.

No, dynamic shapes are not yet supported

Very sad. I would have find another framework (or switch to Julia) for my task.

in practice, this is not so trivial to implement!

Indeed, what I meant is that there is no enrtropy around index i, its advancements are determenistic and predictible, therefore the program depends on it quite straightforwardly. However, the generic case can be much more involved.

[jakevdp]

Note that if you're interested in jit-compiling a function like this, there's no reason you can't do it with straightforward for loops; for example:

@jax.jit
def f(var):
  for i in range(5):
    var = var.at[i, 0].update(jnp.ones([2, i + 1]))
  return var

In this case, all the shapes are static so JAX can handle it just fine. If you want to use scan for this, it's often possible to express your function in terms of padded arrays – let me know if you'd like to see an example of that.

[awav]

Hm, how about performance in this case? What does it mean for the compilation time and memory consumption? Will jax unroll iterations and pre-allocate tensors for all steps?

scan doesn't really work for me unfortunately. The var in my example is actually a big tensor, making copies of it would be really bad for the performance.

[jakevdp]

If you use the for loop approach, JAX will unroll the loops. This can lead to good runtime performance, but if the loops are large, it can make compilation quite expensive.

The alternative is to use scan over padded arrays. Depending on what you're doing, this does not necessarily mean you'll be making copies of your tensor – without more context I'm not sure how to address your concern. Is the code snippet in your question anything close to your actual problem, or is it just a simple toy example to demonstrate how dynamic shapes might arise in scan?

[awav]

A "simple" example would be implementation of the Cholesky algorithm:

1. It has in-place updates,
2. The shape of updates changes from iteration to iteration.

I don't think that it is possible to implement Cholesky with scan, one must have a state and pass it around between iterations.