How do we test out vmap over piles?

[KeAWang]

I'd like to test out the jax.interpreters.batching.Pile feature which I know is still a work in progress. Copying from the dynamic api unittests, I tried

import jax
import jax.numpy as jnp
from jax.interpreters import batching

jax.config.update("jax_array", False)
jax.config.update("jax_dynamic_shapes", True)
jax.config.update("jax_numpy_rank_promotion", "allow")

xs = jax.vmap(lambda n: jnp.arange(n).sum())(jnp.array([3, 1, 4]))

sizes = jnp.array([3, 1, 4])
p1 = jax.vmap(lambda n: jnp.ones((7, n)), out_axes=batching.pile_axis
              )(sizes)
p2 = jax.vmap(lambda n: jnp.ones((n, 7)), out_axes=batching.pile_axis
              )(sizes)
y = jax.vmap(jnp.dot, in_axes=batching.pile_axis, out_axes=0,
             axis_size=3)(p1, p2)

def pile_map(f):
    def mapped(*piles):
        return jax.vmap(f, in_axes=batching.pile_axis, out_axes=batching.pile_axis,
                        axis_size=piles[0].aval.length)(*piles)
    return mapped

p = jax.vmap(jnp.arange, out_axes=batching.pile_axis)(jnp.array([3, 1, 4]))
y = pile_map(jnp.dot)(p, p)

but I keep getting an error

File ~/miniconda3/envs/dev/lib/python3.10/site-packages/jax/_src/pjit.py:873, in <genexpr>(.0)
    [869](file:///lfs/local/0/alexkw/miniconda3/envs/dev/lib/python3.10/site-packages/jax/_src/pjit.py?line=868)   assert not in_shardings_flat
    [870](file:///lfs/local/0/alexkw/miniconda3/envs/dev/lib/python3.10/site-packages/jax/_src/pjit.py?line=869)   return (), ()
    [872](file:///lfs/local/0/alexkw/miniconda3/envs/dev/lib/python3.10/site-packages/jax/_src/pjit.py?line=871) in_axis_resources_flat = tuple(
--> [873](file:///lfs/local/0/alexkw/miniconda3/envs/dev/lib/python3.10/site-packages/jax/_src/pjit.py?line=872)     i if _is_from_gda(i) or is_auto(i) else i._parsed_pspec
    [874](file:///lfs/local/0/alexkw/miniconda3/envs/dev/lib/python3.10/site-packages/jax/_src/pjit.py?line=873)     for i in in_shardings_flat)
    [876](file:///lfs/local/0/alexkw/miniconda3/envs/dev/lib/python3.10/site-packages/jax/_src/pjit.py?line=875) # This check should be above local_to_global call below otherwise if
    [877](file:///lfs/local/0/alexkw/miniconda3/envs/dev/lib/python3.10/site-packages/jax/_src/pjit.py?line=876) # `FROM_GDA` is passed to any input other than GDA, a ugly error message
    [878](file:///lfs/local/0/alexkw/miniconda3/envs/dev/lib/python3.10/site-packages/jax/_src/pjit.py?line=877) # will be raised because get_array_mapping (in local_to_global) of a
    [879](file:///lfs/local/0/alexkw/miniconda3/envs/dev/lib/python3.10/site-packages/jax/_src/pjit.py?line=878) # FROM_GDA cannot happen.
    [880](file:///lfs/local/0/alexkw/miniconda3/envs/dev/lib/python3.10/site-packages/jax/_src/pjit.py?line=879) tree_map(_check_resources_mismatch, in_axis_resources_flat, is_gda)

AttributeError: 'UnspecifiedValue' object has no attribute '_parsed_pspec'

I'm running this on jax 0.4.6 on a CPU. What's the right way to test this feature out?

[mattjj]

Haha, thanks for finding our hidden features! It's really not ready yet; just a concept prototype.

That said, if you run your script with JAX_JIT_PJIT_API_MERGE=0 python filename.py it should work.

[KeAWang]

Amazing! @slinderman was just telling us today about how long he's been waiting for this :)
Do piles work with regular jit? Just now, I found that if I disable jit I could get my code running.

[mattjj]

They can, and someday they will, but the support may be spotty at the moment. I actually forget the exact status. If you open issues about specific errors we can try to take a look, though pile_map isn't a top priority at the moment, and there are likely so many holes that anything beyond simple demos may not work.

Some things are hard to batch efficiently using the current pile representation. For example, self-attention with ragged sequence lengths! We might have to use custom kernels for some operations. Or we might have to switch to a padded-and-stacked representation instead of a concatenated one.