Repository: google/jax

Issue:
document lax.axis_index to get mapped element index from axis name
Hello JAX team,
First of all, thanks a lot for cool JAX. This is a great tool.

This issue contains a question about your plans regarding `pmap` and a feature request.
What `pmap`ed function receives as an input, is that a slice of input array or a full array? What happens in case of nested `pmap`s? Guess JAX traces a function and runs copies of it on sharded chunks of an input array. If so, that would super useful to have access to an index of that shard, that would allow encoding more complex algorithms. Otherwise, it still would be great to have an access to the shard index.

<details><summary>Example</summary>
<p>

```python
@partial(pmap, axis_name='rows')
@partial(pmap, axis_name='cols')
def f(x):
  # Q:  When the code is running, what the `x.shape` actually is?
  # 
  # Dummy example of how to get shard id:
  # shard_row_id = __global_axis_name__["rows"]
  # shard_col_id = __global_axis_name__["cols"]
  #
  # Do something smart with it! E.g. process tensor in `chess` like style - blacks are positive, whites are negative.
  pass
```

</p>
</details>

Thanks



Discussion:
Thanks for the kind words!

I think you might find answers to these in [the SPMD cookbook](https://colab.sandbox.google.com/github/google/jax/blob/master/cloud_tpu_colabs/Pmap_Cookbook.ipynb). But perhaps you've already looked at that, based on your question, and you want clarification.

> Guess JAX traces a function and runs copies of it on sharded chunks of an input array.

That sounds right to me!

> What pmaped function receives as an input, is that a slice of input array or a full array?

It's a slice, like you'd get from the expression `x[i]` for integer `i` where `x` is a NumPy ndarray (so in particular there's no singleton axis).

In your example, the shape of `x` in the body of `f` depends on the shape of the argument to which the pmapped function `f` is applied. In general, the shape of `x` in the body will be missing two leading axes. That is,

```python
@partial(pmap, axis_name='rows')
@partial(pmap, axis_name='cols')
def f(x):
  print(x.shape)
  ...

f(np.ones((4, 2, 9, 9)))  # prints (9, 9)
```

While `pmap` is most fun when programming multiple GPUs or multiple TPU cores, you can experiment with it locally by setting the environment variable `XLA_FLAGS=--xla_force_host_platform_device_count=8` or something like that. See below for an example.

> What happens in case of nested pmaps?

In terms of what values get computed, the semantics are compositional. Without using collectives, it's just like doing nested `vmap`s, which themselves are roughly defined by `vmap(f)(xs) = np.stack([f(x) for x in xs])`. With collectives, every pmap can bind an axis variable name, and collectives refer to those names. Here's an example from the pmap cookbook:

```python
@partial(pmap, axis_name='rows')
@partial(pmap, axis_name='cols')
def f(x):
  row_normed = x / lax.psum(x, 'rows')
  col_normed = x / lax.psum(x, 'cols')
  doubly_normed = x / lax.psum(x, ('rows', 'cols'))
  return row_normed, col_normed, doubly_normed

x = np.arange(8.).reshape((4, 2))
a, b, c = f(x)
print(a)
print(b)
prrint(c)
```

```
[[0.        , 0.0625    ],
 [0.16666667, 0.1875    ],
 [0.33333333, 0.3125    ],
 [0.5       , 0.4375    ]])
[[0.        , 1.        ],
 [0.4       , 0.6       ],
 [0.44444444, 0.55555556],
 [0.46153846, 0.53846154]])
[[0.        , 0.03571429],
 [0.07142857, 0.10714286],
 [0.14285714, 0.17857143],
 [0.21428571, 0.25      ]])
```

Here are some images from slides that might illustrate things:

![image](https://user-images.githubusercontent.com/1458824/77836167-235e8080-7111-11ea-9516-0572ea3ace28.png)


![image](https://user-images.githubusercontent.com/1458824/77836176-2eb1ac00-7111-11ea-88d2-4d5bce53ab05.png)


![image](https://user-images.githubusercontent.com/1458824/77836185-396c4100-7111-11ea-9f4b-7667c2ce2c81.png)

> If so, that would super useful to have access to an index of that shard, that would allow encoding more complex algorithms.

Check out `lax.axis_index`, which seems to be undocumented right now:

```
$ env XLA_FLAGS=--xla_force_host_platform_device_count=8 ipython                                               (base)
Python 3.7.4 (default, Aug 13 2019, 20:35:49)
Type 'copyright', 'credits' or 'license' for more information
IPython 7.8.0 -- An enhanced Interactive Python. Type '?' for help.

In [1]: from jax import pmap
In [2]: from jax import lax

In [3]: import jax.numpy as np

In [4]: pmap(lambda x: x + lax.axis_index('i'), axis_name='i')(np.zeros(8))
Out[4]: ShardedDeviceArray([0., 1., 2., 3., 4., 5., 6., 7.], dtype=float32)

In [6]: @partial(pmap, axis_name='rows')
   ...: @partial(pmap, axis_name='cols')
   ...: def f(_):
   ...:     return lax.axis_index('rows'), lax.axis_index('cols')
   ...:

In [7]: x, y = f(np.zeros((4, 2)))

In [8]: print(x)
[[0 0]
 [1 1]
 [2 2]
 [3 3]]

In [9]: print(y)
[[0 1]
 [0 1]
 [0 1]
 [0 1]]
```

Is that what you have in mind?
Awesome response! Thank you @mattjj!
The `lax.axis_index` is exactly what I was searching! Please, document this method )))

Resolving patch:
diff --git a/jax/interpreters/pxla.py b/jax/interpreters/pxla.py
--- a/jax/interpreters/pxla.py
+++ b/jax/interpreters/pxla.py
@@ -248,6 +248,43 @@ def apply_parallel_primitive(prim, *args, **params):
 
 
 def axis_index(axis_name):
+  """Return the index along the pmapped axis ``axis_name``.
+
+  Args:
+    axis_name: hashable Python object used to name the pmapped axis (see the
+      ``pmap`` docstring for more details).
+
+  Returns:
+    An integer representing the index.
+
+  For example, with 8 XLA devices available:
+
+  >>> from functools import partial
+  >>> @partial(pmap, axis_name='i')
+  ... def f(_):
+  ...   return lax.axis_index('i')
+  ...
+  >>> f(np.zeros(4))
+  ShardedDeviceArray([0, 1, 2, 3], dtype=int32)
+  >>> f(np.zeros(8))
+  ShardedDeviceArray([0, 1, 2, 3, 4, 5, 6, 7], dtype=int32)
+  >>> @partial(pmap, axis_name='i')
+  ... @partial(pmap, axis_name='j')
+  ... def f(_):
+  ...   return lax.axis_index('i'), lax.axis_index('j')
+  ...
+  >>> x, y = f(np.zeros((4, 2)))
+  >>> print(x)
+  [[0 0]
+   [1 1]
+   [2 2]
+   [3 3]]
+  >>> print(y)
+  [[0 1]
+   [0 1]
+   [0 1]
+   [0 1]]
+  """
   dynamic_axis_env = _thread_local_state.dynamic_axis_env
   frame = dynamic_axis_env[axis_name]
   sizes = dynamic_axis_env.sizes[:dynamic_axis_env.index(frame)+1]


Test patch: