Repository: google/jax

Issue:
scatter_add requires dimension numbers, default value invalid
```python3
jax.lax.scatter_add(np.zeros(10), np.asarray([3, 1]), np.asarray([5., 7.]))
```

fails with `AssertionError` because of `assert type(dimension_numbers) is ScatterDimensionNumbers`

The following works, but there is no documentation on these parameters, so I inferred them from the tests.

```python3
dnums = jax.lax.ScatterDimensionNumbers(update_window_dims=(), inserted_window_dims=(0,), scatter_dims_to_operand_dims=(0,), index_vector_dim=1)
jax.lax.scatter_add(np.zeros(10), np.asarray([3, 1]), np.asarray([5., 7.]), dnums)
```

Also, it might be worth adding `jax.lax.scatter` with the first argument set to zeros.


Discussion:
Thanks for the issue report!

I'll add doc strings for these functions (we haven't gotten around to documenting all of `lax` yet), but essentially gather and scatter_add are direct wrappers around XLA's [gather](https://www.tensorflow.org/xla/operation_semantics#gather) and [scatter](https://www.tensorflow.org/xla/operation_semantics#scatter) primitives, so you can refer to the XLA documentation for details. That said, I personally find their APIs too complicated. Expect us to simplify these APIs in the future.

For most `gather` use cases, it's probably better to use NumPy indexing syntax, which will lower to gather already.

For scatter, unfortunately we don't yet have an equivalent to the gather syntax, but we definitely plan to add some kind of update syntax that looks a lot more like NumPy indexed updates! The closest I can offer you right now is some gradient trickery that exploits the fact that scatter-add is the transpose of gather, and we do have nice indexing support for gather:

```
In [49]: x = np.ones((10,), onp.float32)

In [50]: y = onp.array([2, 4, 5])

In [51]: def scatter_add(f, base, updates):
    ...:     _, grad = jax.vjp(f, base)
    ...:     (out,) = grad(updates)
    ...:     return base + out
    ...:

In [52]: scatter_add(lambda a: a[y], x, np.array([7, 4, 32], np.float32)).copy()
Out[52]: array([ 1.,  1.,  8.,  1.,  5., 33.,  1.,  1.,  1.,  1.], dtype=float32)
```

Expect us to offer something along those lines in the future, but it won't be that exact implementation.

Resolving patch:
diff --git a/jax/lax.py b/jax/lax.py
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -503,12 +503,52 @@ def dynamic_update_slice(operand, update, start_indices):
   return dynamic_update_slice_p.bind(operand, update, start_indices,
                                      update_shape=update.shape)
 
-def gather(operand, start_indices, dimension_numbers=None, slice_sizes=None):
+def gather(operand, start_indices, dimension_numbers, slice_sizes):
+  """Gather operator.
+
+  Wraps `XLA's Gather operator
+  <https://www.tensorflow.org/xla/operation_semantics#gather>`_.
+
+  The semantics of gather are complicated, and its API might change in the
+  future. For most use cases, you should prefer `Numpy-style indexing
+  <https://docs.scipy.org/doc/numpy-1.16.0/reference/arrays.indexing.html>`_
+  (e.g., `x[:, (1,4,7), ...]`), rather than using `gather` directly.
+
+  Args:
+    operand: an array from which slices should be taken
+    start_indices: the indices at which slices should be taken
+    dimension_numbers: a `lax.GatherDimensionNumbers` object that describes
+      how dimensions of `operand`, `start_indices` and the output relate.
+    slice_sizes: the size of each slice. Must be a sequence of non-negative
+      integers with size equal to `ndim(operand)`.
+
+  Returns:
+    An array containing the gather output.
+  """
   return gather_p.bind(
       operand, start_indices, dimension_numbers=dimension_numbers,
       slice_sizes=tuple(slice_sizes), operand_shape=operand.shape)
 
-def scatter_add(operand, scatter_indices, updates, dimension_numbers=None):
+def scatter_add(operand, scatter_indices, updates, dimension_numbers):
+  """Scatter operator.
+
+  Wraps `XLA's Scatter operator
+  <https://www.tensorflow.org/xla/operation_semantics#scatter>`_.
+
+  The semantics of scatter are complicated and its API is subject to change.
+
+  Args:
+    operand: an array to which the scatter should be applied
+    scatter_indices: an array that gives the indices in `operand` to which each
+      update in `updates` should be applied.
+    updates: the updates that should be scattered onto `operand`.
+    dimension_numbers: a `lax.ScatterDimensionNumbers` object that describes
+      how dimensions of `operand`, `start_indices`, `updates` and the output
+      relate.
+
+  Returns:
+    An array containing the sum of `operand` and the scattered updates.
+  """
   jaxpr, consts = _reduction_jaxpr(add, _const(operand, 0))
   return scatter_p.bind(
       operand, scatter_indices, updates, update_jaxpr=jaxpr,
@@ -2364,10 +2404,30 @@ def _dynamic_update_slice_translation_rule(c, operand, update, start_indices,
 
 
 
-GatherDimensionNumbers = collections.namedtuple(
+class GatherDimensionNumbers(collections.namedtuple(
     "GatherDimensionNumbers",
     ["offset_dims", "collapsed_slice_dims", "start_index_map",
-     "index_vector_dim"])
+     "index_vector_dim"])):
+  """
+  Describes the dimension number arguments to an `XLA's Gather operator
+  <https://www.tensorflow.org/xla/operation_semantics#gather>`_. See the XLA
+  documentation for more details of what the dimension numbers mean.
+
+  Args:
+    offset_dims: the set of dimensions in the `gather` output that offset into
+      an array sliced from `operand`. Must be a tuple of integers in ascending
+      order, each representing a dimension number of the output.
+    collapsed_slice_dims: the set of dimensions `i` in `operand` that have
+      `slice_sizes[i] == 1` and that should not have a corresponding dimension
+      in the output of the gather. Must be a tuple of integers in ascending
+      order.
+    start_index_map: for each dimension in `start_indices`, gives the
+      corresponding dimension in `operand` that is to be sliced. Must be a
+      tuple of integers with size equal to `ndim(start_indices)`.
+    index_vector_dim: describes which dimension of `start_indices` "contains"
+      the start indices. If equal to `len(start_indices)` the indices are
+      taken to be scalars.
+  """
 
 def _gather_dimensions_proto(dimension_numbers):
   assert type(dimension_numbers) is GatherDimensionNumbers
@@ -2506,10 +2566,29 @@ def _gather_batching_rule(batched_args, batch_dims, dimension_numbers,
 batching.primitive_batchers[gather_p] = _gather_batching_rule
 
 
-ScatterDimensionNumbers = collections.namedtuple(
+class ScatterDimensionNumbers(collections.namedtuple(
     "ScatterDimensionNumbers",
     ["update_window_dims", "inserted_window_dims",
-     "scatter_dims_to_operand_dims", "index_vector_dim"])
+     "scatter_dims_to_operand_dims", "index_vector_dim"])):
+  """
+  Describes the dimension number arguments to an `XLA's Scatter operator
+  <https://www.tensorflow.org/xla/operation_semantics#scatter>`_. See the XLA
+  documentation for more details of what the dimension numbers mean.
+
+  Args:
+    update_window_dims: the set of dimensions in the `updates` that are window
+      dimensions. Must be a tuple of integers in ascending
+      order, each representing a dimension number.
+    inserted_window_dims: the set of size 1 window dimensions that must be inserted
+      into the shape of `updates`. Must be a tuple of integers in ascending
+      order, each representing a dimension number of the output. These are the
+      mirror image of `collapsed_slice_dims` in the case of `gather`.
+    scatter_dims_to_operand_dims: for each dimension in `scatter_indices`, gives
+      the corresponding dimension in `operand`.
+    index_vector_dim: describes which dimension of `scatter_indices` "contains"
+      the start indices. If equal to `len(scatter_indices)` the indices are
+      taken to be scalars.
+  """
 
 def _scatter_dimensions_proto(dimension_numbers):
   assert type(dimension_numbers) is ScatterDimensionNumbers
@@ -2523,7 +2602,7 @@ def _scatter_dimensions_proto(dimension_numbers):
 
 def _scatter_dtype_rule(operand, scatter_indices, updates, **kwargs):
   if not onp.issubdtype(scatter_indices.dtype, onp.integer):
-    raise ValueError("start_indices must have an integer type")
+    raise ValueError("scatter_indices must have an integer type")
   _check_same_dtypes("scatter", False, operand.dtype, updates.dtype)
   return xla_bridge.canonicalize_dtype(operand.dtype)
 


Test patch: