Cannot generate arrays with dimensions set at runtime

[joacorapela]

I want to create in JAX matrices and vectors with dimensions set at runtime. Can this be done in JAX?

If the dimension of these matrices and vectors is fixed at runtime, I can do it with the following code

import jax
import jax.numpy as jnp

jax.config.update("jax_enable_x64", True)

def build_data(key, N, dim):
    def body_fun(i, val):
        chols, vecs, key = val
        key, subkey = jax.random.split(key)
        chol = jax.random.uniform(subkey,
                                  shape=(dim, dim),
                                  dtype=jnp.double)
        key, subkey = jax.random.split(key)
        vec = jax.random.uniform(subkey, shape=(dim,), dtype=jnp.double)
        chols = chols.at[i, :, :].set(chol)
        vecs = vecs.at[i, :].set(vec)
        new_val = (chols, vecs, key)
        return new_val
    chols = jnp.empty(shape=(N, dim, dim), dtype=jnp.double)
    vecs = jnp.empty(shape=(N, dim), dtype=jnp.double)
    val = (chols, vecs, key)
    val = jax.lax.fori_loop(0, N, body_fun, val)
    return val

key = jax.random.PRNGKey(324234)
dim = 4
N = 5
dims_options = jnp.arange(5, 10)
chols, vecs, key = build_data(key=key, N=N, dim=dim)

However, if their dimensions is determined at runtime, I cannot create similar random
vectors and matrices with the following code.

import jax
import jax.numpy as jnp

jax.config.update("jax_enable_x64", True)

def build_data(key, dims):
    def body_fun(i, val):
        chols, vecs, dims, key = val
        key, subkey = jax.random.split(key)
        chol = jax.random.uniform(subkey,
                                  shape=(dims.at[i].get(), dims.at[i].get()),
                                  dtype=jnp.double)
        key, subkey = jax.random.split(key)
        vec = jax.random.uniform(subkey, shape=(dims.at[i].get(),), dtype=jnp.double)
        chols.append(chol)
        vecs.append(vec)
        new_val = (chols, vecs, dims, key)
        return new_val
    chols = []
    vecs = []
    val = (chols, vecs, dims, key)
    val = jax.lax.fori_loop(0, N, body_fun, val)
    return val

N = 5
dims_options = jnp.arange(5, 10)

key = jax.random.PRNGKey(324234)
key, subkey = jax.random.split(key)
dims = jax.random.choice(key=subkey, a=dims_options, shape=(N,))

chols, vecs, key = build_data(key=key, dims=dims)

The previous code has at least two problems:

1. The following line generates the error below

           chol = jax.random.uniform(subkey,
                                     shape=(dims.at[i].get(), dims.at[i].get()),
                                     dtype=jnp.double)

   TypeError: Shapes must be 1D sequences of concrete values of integer type, got
   (Traced<ShapedArray(int64[])>with<DynamicJaxprTrace(level=1/0)>,
   Traced<ShapedArray(int64[])>with<DynamicJaxprTrace(level=1/0)>).

2. I cannot use expanding lists in the carry variable.

Is there a workaround in JAX to create matrices with dimensions set at runtime?

Joaquin

[joacorapela]

It seems that this cannot be done with fori_loop so I did it with a for loop i #19576 However I have a related fori_loop problem in #19576.

[jakevdp]

No, you cannot create arrays with shapes set at runtime, unless you run in eager mode (e.g. outside jit and other transformations) See JAX Sharp Bits: Dynamic Shapes for more information.

Typical workarounds are to set static shapes of some maximum size, and use padded representations for dynamically-shaped arrays within your program.