vmap a function with variable length for loop

[JacobHelwig]

I have an array X of dimension [N, dim1,...,dimK] and would like to perform some calculations along the first axis, i.e., on N arrays of dimension [dim1,...,dimK]. Each of these N calculations involves a for loop whose length is known but varies for each of the N entries.

I would like to vmap these calculations along the first axis. Here is an example of what I am trying to do with the expected output using a for loop in place of vmap:

>>> import jax
>>> import jax.numpy as jnp
>>> n = 3
>>> dim = 2
>>> X = jnp.arange(start=1, stop=n + 1).repeat(dim**2).reshape(n, dim, dim) * jnp.pi

>>> reps = (jnp.arange(n) + 1) * 2
>>> X
DeviceArray([[[3.1415927, 3.1415927],
              [3.1415927, 3.1415927]],

             [[6.2831855, 6.2831855],
              [6.2831855, 6.2831855]],

             [[9.424778 , 9.424778 ],
              [9.424778 , 9.424778 ]]], dtype=float32, weak_type=True)
>>> reps
DeviceArray([2, 4, 6], dtype=int32)
>>> def calculations(Z, n_reps):
...     sum0 = Z.sum()
...     jax.debug.print("Printing {sum0} {n_reps} times", sum0=sum0, n_reps=n_reps)
...     def printer(idx, val):
...         jax.debug.print("{}", val)
...         return val
...     jax.lax.fori_loop(lower=0, upper=n_reps, body_fun=printer, init_val=sum0)
... 
>>> for i in range(n):
...     calculations(X[i], reps[i])
... 
Printing 12.566370964050293 2 times
12.566370964050293
12.566370964050293
Printing 25.132741928100586 4 times
25.132741928100586
25.132741928100586
25.132741928100586
25.132741928100586
Printing 37.69911193847656 6 times
37.69911193847656
37.69911193847656
37.69911193847656
37.69911193847656
37.69911193847656
37.69911193847656

When I try to vmap, the loop is not applied the correct number of times:

>>> jax.vmap(calculations)(X, reps)
Printing 12.566370964050293 2 times
Printing 25.132741928100586 4 times
Printing 37.69911193847656 6 times
12.566370964050293
25.132741928100586
37.69911193847656
12.566370964050293
25.132741928100586
37.69911193847656
12.566370964050293
25.132741928100586
37.69911193847656
12.566370964050293
25.132741928100586
37.69911193847656
12.566370964050293
25.132741928100586
37.69911193847656
12.566370964050293
25.132741928100586
37.69911193847656

Any suggestions would be appreciated

[sharadmv]

tl;dr - this an example of why we call it "debug" print and is working as intended.

"In the name of debugging, jax.debug.print can reveal information about how computations are evaluated"

This is a really great example of how jax.debug.print exposes how JAX is actually running your computation.

When you do vmap-of-fori_loop where the number of steps is different across the mapped axis (as is in your example), we convert it into what we call a "padded while loop". A "padded while loop" runs the batched body until every batch member has finished, which involves keeping track of which batch member has finished in an array of bools and then ignoring the outputs of the batched body for each of the members that had completed already. It looks roughly like this:

def batched_body(finished_and_carry):
  finished, batched_carry = finished_and_carry
  batched_carry_out = vmap(body)(batched_carry)
  batched_carry = jnp.where(finished, batched_carry, batched_carry_out)
  finished = vmap(cond)(batched_carry)
  return finished, batched_carry

Note that this transformation preserves correctness if the body and cond functions are pure!. If they have side-effects, we will be running the side-effects even if the batch member has completed the loop.

Ordinarily, we disallow side-effects inside of batched while loops like this (for example, checkify.check is not allowed inside of a batched while loop). However, we explicitly allow jax.debug.print/breakpoint/callback in this setting because their purpose is to reveal what is happening in JAX computations.