Repository: google/jax

Issue:
vmap of dynamic_slice of a scalar fails
Specifically, doing
```python
In [1]: import jax.numpy as np

In [2]: from jax import lax

In [3]: import jax

In [4]: jax.vmap(lambda x: lax.dynamic_slice(x, [], []))(np.array([1, 2, 3, 4]))
```
gives the following error
```text
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-4-ae0083eeb242> in <module>
----> 1 jax.vmap(lambda x: lax.dynamic_slice(x, [], []))(np.array([1, 2, 3, 4]))

~/dev/jax/jax/api.py in batched_fun(*args)
    878     _ = _mapped_axis_size(in_tree, args_flat, in_axes_flat, "vmap")
    879     out_flat = batching.batch(flat_fun, args_flat, in_axes_flat,
--> 880                               lambda: flatten_axes("vmap out_axes", out_tree(),
    881                                                    out_axes))
    882     return tree_unflatten(out_tree(), out_flat)

~/dev/jax/jax/interpreters/batching.py in batch(fun, in_vals, in_dims, out_dim_dests)
     32   # executes a batched version of `fun` following out_dim_dests
     33   batched_fun = batch_fun(fun, in_dims, out_dim_dests)
---> 34   return batched_fun.call_wrapped(*in_vals)
     35
     36 @lu.transformation_with_aux

~/dev/jax/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    148     gen = None
    149
--> 150     ans = self.f(*args, **dict(self.params, **kwargs))
    151     del args
    152     while stack:

<ipython-input-4-ae0083eeb242> in <lambda>(x)
----> 1 jax.vmap(lambda x: lax.dynamic_slice(x, [], []))(np.array([1, 2, 3, 4]))

~/dev/jax/jax/lax/lax.py in dynamic_slice(operand, start_indices, slice_sizes)
    734   start_indices = _dynamic_slice_indices(operand, start_indices)
    735   return dynamic_slice_p.bind(operand, *start_indices,
--> 736                               slice_sizes=tuple(slice_sizes))
    737
    738 def dynamic_update_slice(operand: Array, update: Array,

~/dev/jax/jax/core.py in bind(self, *args, **kwargs)
    276
    277     tracers = map(top_trace.full_raise, args)
--> 278     out_tracer = top_trace.process_primitive(self, tracers, kwargs)
    279     if self.multiple_results:
    280       return map(full_lower, out_tracer)

~/dev/jax/jax/interpreters/batching.py in process_primitive(self, primitive, tracers, params)
    132       # TODO(mattjj,phawkins): if no rule implemented, could vmap-via-map here
    133       batched_primitive = get_primitive_batcher(primitive)
--> 134       val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
    135       if primitive.multiple_results:
    136         return map(partial(BatchTracer, self), val_out, dim_out)

~/dev/jax/jax/lax/lax.py in _dynamic_slice_batching_rule(batched_args, batch_dims, slice_sizes)
   3696   dnums = GatherDimensionNumbers(offset_dims=dims, collapsed_slice_dims=(),
   3697                                  start_index_map=dims)
-> 3698   index, index_bdim = _batch_dynamic_slice_indices(start_indices, start_idx_bds)
   3699   return _gather_batching_rule(
   3700     [operand, index], [operand_bd, index_bdim], dimension_numbers=dnums,

~/dev/jax/jax/lax/lax.py in _batch_dynamic_slice_indices(indices, bdims)
   3676   size = next((x.shape[i] for x, i in zip(indices, bdims) if i is not None), -1)
   3677   if size < 0:
-> 3678     return concatenate([broadcast(i, (1,)) for i in indices], 0), None
   3679   indices = concatenate(
   3680     [broadcast_in_dim(x, (size, 1),

~/dev/jax/jax/lax/lax.py in concatenate(operands, dimension)
    446     An array containing the concatenation.
    447   """
--> 448   return concatenate_p.bind(*operands, dimension=dimension)
    449
    450 Precision = xla_client.PrecisionConfig.Precision

~/dev/jax/jax/core.py in bind(self, *args, **kwargs)
    273     top_trace = find_top_trace(args)
    274     if top_trace is None:
--> 275       return self.impl(*args, **kwargs)
    276
    277     tracers = map(top_trace.full_raise, args)

~/dev/jax/jax/interpreters/xla.py in apply_primitive(prim, *args, **params)
    222 def apply_primitive(prim, *args, **params):
    223   """Impl rule that compiles and runs a single primitive 'prim' using XLA."""
--> 224   compiled_fun = xla_primitive_callable(prim, *unsafe_map(arg_spec, args), **params)
    225   return compiled_fun(*args)
    226

~/dev/jax/jax/interpreters/xla.py in xla_primitive_callable(prim, *arg_specs, **params)
    238     return _xla_callable(lu.wrap_init(prim_fun), device, None, "prim", donated_invars,
    239                          *arg_specs)
--> 240   aval_out = prim.abstract_eval(*avals, **params)
    241   if not prim.multiple_results:
    242     handle_result = aval_to_result_handler(device, aval_out)

~/dev/jax/jax/lax/lax.py in standard_abstract_eval(prim, shape_rule, dtype_rule, *args, **kwargs)
   1845   assert all(isinstance(arg, UnshapedArray) for arg in args), args
   1846   least_specialized = _max(
-> 1847       map(type, args), key=operator.attrgetter('array_abstraction_level'))
   1848   if least_specialized is ConcreteArray:
   1849     return ConcreteArray(prim.impl(*[x.val for x in args], **kwargs))

ValueError: max() arg is an empty sequence
```
I'm just writing a test for this now and will have a go at fixing it.


Discussion:

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -3673,6 +3673,8 @@ def _dynamic_slice_transpose_rule(t, operand, *start_indices, slice_sizes):
           [None] * len(start_indices))
 
 def _batch_dynamic_slice_indices(indices, bdims):
+  if len(indices) == 0:
+    return np.array([], 'int32'), None
   size = next((x.shape[i] for x, i in zip(indices, bdims) if i is not None), -1)
   if size < 0:
     return concatenate([broadcast(i, (1,)) for i in indices], 0), None
@@ -3769,8 +3771,8 @@ def _dynamic_update_slice_batching_rule(batched_args, batch_dims):
   # scatter always.
   operand, update, *start_idx = batched_args
   operand_bd, update_bd, *start_idx_bd = batch_dims
-  update_shape = (update.shape if update_bd is batching.not_mapped
-                  else tuple(np.delete(update.shape, update_bd)))
+  update_shape = (np.shape(update) if update_bd is batching.not_mapped
+                  else tuple(np.delete(np.shape(update), update_bd)))
   dims = tuple(range(len(update_shape)))
   dnums = ScatterDimensionNumbers(update_window_dims=dims,
                                   inserted_window_dims=(),


Test patch:
diff --git a/tests/batching_test.py b/tests/batching_test.py
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -960,6 +960,22 @@ def f(index1, index2):
     expected = g(np.asarray([1]), np.asarray([2]))
     self.assertAllClose(ans, expected)
 
+  def testIssue3883(self):
+    def scalar_f(x):
+      return lax.dynamic_slice(x, [], [])
+
+    xs = jnp.array([1, 2, 3, 4])
+    ans = vmap(scalar_f)(xs)
+    expected = jnp.array([scalar_f(x) for x in xs])
+    self.assertAllClose(ans, expected)
+
+    def scalar_f2(x):
+      return lax.dynamic_update_slice(x, 7, [])
+
+    xs = jnp.array([1, 2, 3, 4])
+    ans = vmap(scalar_f2)(xs)
+    expected = jnp.array([scalar_f2(x) for x in xs])
+    self.assertAllClose(ans, expected)
 
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())