Remove jnp.atleast_1d in subsetting

Author: michalk8

pyproject.toml

@@ -120,7 +120,10 @@ markers = [
     "fast: Mark tests as fast.",
 ]
 filterwarnings = [
+    "ignore:\\n*.*scipy.sparse array",
     "ignore:jax.random.KeyArray is deprecated:DeprecationWarning",
+    "ignore:.*jax.config:DeprecationWarning",
+    "ignore:jax.core.Shape is deprecated:DeprecationWarning:chex",
 ]
 
 [tool.coverage.run]

src/ott/geometry/geometry.py

@@ -740,9 +740,9 @@ def subset_fn(
       if arr is None:
         return None
       if src_ixs is not None:
-        arr = arr[jnp.atleast_1d(src_ixs)]
+        arr = arr[src_ixs, ...]
       if tgt_ixs is not None:
-        arr = arr[:, jnp.atleast_1d(tgt_ixs)]
+        arr = arr[:, tgt_ixs]
       return arr  # noqa: RET504
 
     return self._mask_subset_helper(
@@ -877,12 +877,11 @@ def _m_normed_ones(self) -> jnp.ndarray:
     return arr / jnp.sum(arr)
 
   @staticmethod
-  def _normalize_mask(mask: Optional[Union[int, jnp.ndarray]],
+  def _normalize_mask(mask: Optional[jnp.ndarray],
                       size: int) -> Optional[jnp.ndarray]:
     """Convert array of indices to a boolean mask."""
     if mask is None:
       return None
-    mask = jnp.atleast_1d(mask)
     if not jnp.issubdtype(mask, (bool, jnp.bool_)):
       mask = jnp.isin(jnp.arange(size), mask)
     assert mask.shape == (size,)

src/ott/geometry/low_rank.py

@@ -251,7 +251,7 @@ def subset_fn(
         arr: Optional[jnp.ndarray],
         ixs: Optional[jnp.ndarray],
     ) -> jnp.ndarray:
-      return arr if arr is None or ixs is None else arr[jnp.atleast_1d(ixs)]
+      return arr if arr is None or ixs is None else arr[ixs, ...]
 
     return self._mask_subset_helper(
         src_ixs, tgt_ixs, fn=subset_fn, propagate_mask=True, **kwargs

src/ott/geometry/pointcloud.py

@@ -648,7 +648,7 @@ def subset_fn(
         arr: Optional[jnp.ndarray],
         ixs: Optional[jnp.ndarray],
     ) -> jnp.ndarray:
-      return arr if arr is None or ixs is None else arr[jnp.atleast_1d(ixs)]
+      return arr if arr is None or ixs is None else arr[ixs, ...]
 
     return self._mask_subset_helper(
         src_ixs, tgt_ixs, fn=subset_fn, propagate_mask=True, **kwargs

src/ott/utils.py

@@ -40,7 +40,7 @@
 def register_pytree_node(cls: type) -> type:
   """Register dataclasses as pytree_nodes."""
   cls = dataclasses.dataclass()(cls)
-  flatten = lambda obj: jax.tree_flatten(dataclasses.asdict(obj))
+  flatten = lambda obj: jax.tree_util.tree_flatten(dataclasses.asdict(obj))
   unflatten = lambda d, children: cls(**d.unflatten(children))
   jax.tree_util.register_pytree_node(cls, flatten, unflatten)
   return cls

tests/geometry/subsetting_test.py

@@ -59,7 +59,7 @@ def geom_masked(request, pc_masked) -> Tuple[Geom_t, pointcloud.PointCloud]:
 @pytest.mark.fast()
 class TestMaskPointCloud:
 
-  @pytest.mark.parametrize("tgt_ixs", [7, jnp.arange(5)])
+  @pytest.mark.parametrize("tgt_ixs", [[1], jnp.arange(5)])
   @pytest.mark.parametrize("src_ixs", [None, (3, 3)])
   @pytest.mark.parametrize(
       "clazz", [geometry.Geometry, pointcloud.PointCloud, low_rank.LRCGeometry]
@@ -78,12 +78,8 @@ def test_mask(
       geom = clazz(cost_matrix=x @ y.T, scale_cost="mean")
     else:
       geom = clazz(x, y, scale_cost="max_cost", batch_size=5)
-    n = geom.shape[0] if src_ixs is None else 1 if isinstance(
-        src_ixs, int
-    ) else len(src_ixs)
-    m = geom.shape[1] if tgt_ixs is None else 1 if isinstance(
-        tgt_ixs, int
-    ) else len(tgt_ixs)
+    n = geom.shape[0] if src_ixs is None else len(src_ixs)
+    m = geom.shape[1] if tgt_ixs is None else len(tgt_ixs)
 
     if clazz is geometry.Geometry:
       geom_sub = geom.subset(src_ixs, tgt_ixs)