How to dump stablehlo for pmap?

[Young768]

I checked that both jax.jit and jax.pmap return a function that is compiled.
I was trying to get the hlo for paxml, here.

I added the line:

lowered = jax.pmap(_wrapped_step_fn).lower((0,), axis_name=base_layer.PMAP_PARALLEL_AXIS_NAME,)

But I hit error:

ValueError: pmap was requested to map its argument along axis 0, which implies that its rank should be at least 1, but is only 0 (its shape is ())

[mattjj]

Thanks for the question!

You need to give jax.pmap(f).lower arguments that are like the same arguments you'd give to jax.pmap(f) itself. So in particular, the argument (0,) can't be mapped over, hence the error. Also, you need to pass in_axes and axis_name values to jax.pmap as usual, not to .lower.

Maybe you want something more like this?

example_x = jnp.zeros(8)
lowered = jax.pmap(_wrapped_step_fun, in_axes=(0,), axis_name=...).lower(example_x)