Cannot replace a simple for loop by a fori_loop

[joacorapela]

I have several matrices Ai and vectors bi of different dimensions and a simple for loop that solves the linear systems Ai xi = bi. I want to replace the for loop by a fori_loop, but I cannot do it with the code below. Can this be done?

The following example calculates solves_for but when trying to calculate solves_fori_loop I get the error below.

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

def buildData(dims):
    N = len(dims)
    vecs = []
    chols = []
    for i in range(N):
        vec = np.random.uniform(size=(dims[i],))
        chol = np.random.uniform(size=(dims[i], dims[i]))
        vecs.append(jax.device_put(vec))
        chols.append(jax.device_put(chol))
    return vecs, chols

N = 1000
dims_options = jnp.arange(5, 30)
key = jax.random.PRNGKey(324234)
key, subkey = jax.random.split(key)
dims = jax.random.choice(key=subkey, a=dims_options, shape=(N,))

vecs, chols = buildData(dims=dims)
solves_for = [jax.scipy.linalg.cho_solve((chols[k], True), vecs[k])
              for k in range(N)]

def loop_step(k, carry):
    solves, vecs, chols = carry
    solves.append(jax.scipy.linalg.cho_solve((chols[k], True), vecs[k]))
    new_carry = (solves, vecs, chols)
    return new_carry

solves = []
carry0 = (solves, vecs, chols)
solves_fori_loop = jax.lax.fori_loop(0, N, loop_step, carry0)

Error message:

TracerIntegerConversionError Traceback (most recent call last)
File ~/dev/learning/jax/doProblemWithFor.py:34
32 solves = []
33 carry0 = (solves, vecs, chols)
---> 34 solves_fori_loop = jax.lax.fori_loop(0, N, loop_step, carry0)
36 breakpoint()

[... skipping hidden 12 frame]                                                                                                                                                                                                                   

File ~/dev/learning/jax/doProblemWithFor.py:28, in loop_step(k, carry)
26 def loop_step(k, carry):
27 solves, vecs, chols = carry
---> 28 solves.append(jax.scipy.linalg.cho_solve((chols[k], True), vecs[k]))
29 new_carry = (solves, vecs, chols)
30 return new_carry

[... skipping hidden 1 frame]                                                                                                                                                                                                                    

File ~/anaconda3/envs/svGPFAjax/lib/python3.9/site-packages/jax/_src/core.py:1513, in concretization_function_error..error(self, arg)
1512 def error(self, arg):
-> 1513 raise TracerIntegerConversionError(arg)

TracerIntegerConversionError: The index() method was called on traced array with shape int32[].
The error occurred while tracing the function scanned_fun at /nfs/ghome/live/rapela/anaconda3/envs/svGPFAjax/lib/python3.9/site-packages/jax/_src/lax/control_flow/loops.py:1853 for scan. This concrete value was not available in Python because it depends on the value of the argument loop_carry[0].
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerIntegerConversionError
(svGPFAjax) rapela@u237c:~/dev/learning/jax$ ipython doProblemWithFor.py

TracerIntegerConversionError Traceback (most recent call last)
File ~/dev/learning/jax/doProblemWithFor.py:34
32 solves = []
33 carry0 = (solves, vecs, chols)
---> 34 solves_fori_loop = jax.lax.fori_loop(0, N, loop_step, carry0)
36 breakpoint()

[... skipping hidden 12 frame]                                                                                                                                                                                                                   

File ~/dev/learning/jax/doProblemWithFor.py:28, in loop_step(k, carry)
26 def loop_step(k, carry):
27 solves, vecs, chols = carry
---> 28 solves.append(jax.scipy.linalg.cho_solve((chols[k], True), vecs[k]))
29 new_carry = (solves, vecs, chols)
30 return new_carry

[... skipping hidden 1 frame]                                                                                                                                                                                                                    

File ~/anaconda3/envs/svGPFAjax/lib/python3.9/site-packages/jax/_src/core.py:1513, in concretization_function_error..error(self, arg)
1512 def error(self, arg):
-> 1513 raise TracerIntegerConversionError(arg)

TracerIntegerConversionError: The index() method was called on traced array with shape int32[].
The error occurred while tracing the function scanned_fun at /nfs/ghome/live/rapela/anaconda3/envs/svGPFAjax/lib/python3.9/site-packages/jax/_src/lax/control_flow/loops.py:1853 for scan. This concrete value was not available in Python because it depends on the value of the argument loop_carry[0].
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerIntegerConversionError
~

[jakevdp]

Hi - this seems related to your other question at #19561. The problem is that your code attempts to create dynamically-shaped arrays, and these are not supported within jax.jit and other JAX transformations. See JAX Sharp Bits: Dynamic Shapes for more information.