[sparse] bcoo_dynamic_slice: remove unnecessary padding from output

jax/experimental/sparse/bcoo.py

@@ -1964,7 +1964,7 @@ def bcoo_slice(mat, *, start_indices: Sequence[int], limit_indices: Sequence[int
 
   return BCOO((new_data, new_indices), shape=new_shape)
 
-def bcoo_dynamic_slice(mat, start_indices: Sequence[Any], slice_sizes: Sequence[int]):
+def bcoo_dynamic_slice(mat: BCOO, start_indices: Sequence[Any], slice_sizes: Sequence[int]) -> BCOO:
   """Sparse implementation of {func}`jax.lax.dynamic_slice`.
 
   Args:
@@ -2028,11 +2028,16 @@ def bcoo_dynamic_slice(mat, start_indices: Sequence[Any], slice_sizes: Sequence[
     sparse_shape = jnp.expand_dims(sparse_shape, range(mat.n_batch + 1))
 
     keep = jnp.all((new_indices >= starts) & (new_indices < starts + sizes), -1, keepdims=True)
-    new_indices = jnp.where(keep, new_indices - starts, sparse_shape)
+    new_indices = jnp.where(keep, new_indices - starts, sizes)
 
     keep_data = lax.expand_dims(keep[..., 0], range(mat.n_batch + 1, mat.n_batch + 1 + mat.n_dense))
     new_data = jnp.where(keep_data, new_data, 0)
 
+    if mat.nse > np.prod(size_sparse):
+      new_nse = np.prod(size_sparse)
+      new_data, new_indices = _bcoo_sum_duplicates(
+        new_data, new_indices, spinfo=BCOOInfo(shape=new_shape), nse=new_nse)
+
   return BCOO((new_data, new_indices), shape=new_shape)
 
 

tests/sparse_test.py

@@ -993,6 +993,17 @@ def test_bcoo_dynamic_slice(self, shape, dtype, n_batch, n_dense):
     sparse_result = sparse.bcoo_dynamic_slice(Msp, **kwds)
     sparse_result_jit = partial(sparse.bcoo_dynamic_slice, slice_sizes=slice_sizes)(Msp, start_indices)
 
+    # Array layout is the same
+    self.assertEqual(sparse_result.n_batch, Msp.n_batch)
+    self.assertEqual(sparse_result.n_sparse, Msp.n_sparse)
+    self.assertEqual(sparse_result.n_dense, Msp.n_dense)
+
+    # Unnecessary padding eliminated
+    max_nse = np.prod(sparse_result.shape[Msp.n_batch: Msp.n_batch + Msp.n_sparse])
+    self.assertLessEqual(sparse_result.nse, max_nse)
+    self.assertLessEqual(sparse_result_jit.nse, max_nse)
+
+    # Result matches dense computation
     self.assertArraysEqual(dense_result, sparse_result.todense())
     self.assertArraysEqual(dense_result, sparse_result_jit.todense())