Finalize assigning output dimension names in index operation

Author: ricardoV94

pytensor/xtensor/indexing.py

@@ -23,11 +23,28 @@ def as_idx_variable(idx):
         idx = make_slice(idx)
     elif isinstance(idx, Variable) and isinstance(idx.type, SliceType):
         pass
-    elif isinstance(idx, tuple) and len(idx) == 2 and isinstance(idx[0], str):
+    elif (
+        isinstance(idx, tuple)
+        and len(idx) == 2
+        and (
+            isinstance(idx[0], str)
+            or (
+                isinstance(idx[0], tuple | list)
+                and all(isinstance(d, str) for d in idx[0])
+            )
+        )
+    ):
         # Special case for ("x", array) that xarray supports
-        # TODO: Check if this can be used to rename existing xarray dimensions or only for numpy
         dim, idx = idx
-        idx = xtensor_from_tensor(as_tensor(idx), dims=(dim,))
+        if isinstance(idx.type, XTensorType):
+            raise TypeError(
+                "Giving a dimension name to an XTensorVariable indexer is not supported"
+            )
+        if isinstance(dim, str):
+            dims = (dim,)
+        else:
+            dims = tuple(dim)
+        idx = xtensor_from_tensor(as_tensor(idx), dims=dims)
     else:
         # Must be integer indices, we already counted for None and slices
         try:

tests/xtensor/test_indexing.py

@@ -52,7 +52,7 @@ def test_single_adv_indexing_on_existing_dim():
     idx_test = np.array([0, 1, 0, 2], dtype=int)
     xidx_test = DataArray(idx_test, dims=("a",))
 
-    # Three equivalent ways of indexing a->a
+    # Equivalent ways of indexing a->a
     y = x[idx]
     fn = xr_function([x, idx], y)
     res = fn(x_test, idx_test)
@@ -65,6 +65,12 @@ def test_single_adv_indexing_on_existing_dim():
     expected_res = x_test[(("a", idx_test),)]
     xr_assert_allclose(res, expected_res)
 
+    y = x[((("a",), idx),)]
+    fn = xr_function([x, idx], y)
+    res = fn(x_test, idx_test)
+    expected_res = x_test[((("a",), idx_test),)]
+    xr_assert_allclose(res, expected_res)
+
     y = x[xidx]
     fn = xr_function([x, xidx], y)
     res = fn(x_test, xidx_test)
@@ -81,13 +87,19 @@ def test_single_vector_indexing_on_new_dim():
     idx_test = np.array([0, 1, 0, 2], dtype=int)
     xidx_test = DataArray(idx_test, dims=("a",))
 
-    # Two equivalent ways of indexing a->new_a
+    # Equivalent ways of indexing a->new_a
     y = x[(("new_a", idx),)]
     fn = xr_function([x, idx], y)
     res = fn(x_test, idx_test)
     expected_res = x_test[(("new_a", idx_test),)]
     xr_assert_allclose(res, expected_res)
 
+    y = x[((["new_a"], idx),)]
+    fn = xr_function([x, idx], y)
+    res = fn(x_test, idx_test)
+    expected_res = x_test[((["new_a"], idx_test),)]
+    xr_assert_allclose(res, expected_res)
+
     y = x[xidx.rename(a="new_a")]
     fn = xr_function([x, xidx], y)
     res = fn(x_test, xidx_test)
@@ -176,6 +188,34 @@ def test_matrix_indexing():
     xr_assert_allclose(res, expected_res)
 
 
+def test_assign_multiple_out_dims():
+    x = xtensor("x", shape=(5, 7), dims=("a", "b"))
+    idx1 = tensor("idx1", dtype=int, shape=(4, 3))
+    idx2 = tensor("idx2", dtype=int, shape=(3, 2))
+    out = x[(("out1", "out2"), idx1), (["out2", "out3"], idx2)]
+
+    fn = xr_function([x, idx1, idx2], out)
+
+    rng = np.random.default_rng()
+    x_test = xr_arange_like(x)
+    idx1_test = rng.binomial(n=4, p=0.5, size=(4, 3))
+    idx2_test = rng.binomial(n=4, p=0.5, size=(3, 2))
+    res = fn(x_test, idx1_test, idx2_test)
+    expected_res = x_test[(("out1", "out2"), idx1_test), (["out2", "out3"], idx2_test)]
+    xr_assert_allclose(res, expected_res)
+
+
+def test_assign_dims_xtensor_fails():
+    x = xtensor("x", shape=(5, 7), dims=("a", "b"))
+    idx1 = xtensor("idx1", dtype=int, shape=(4,), dims=("c",))
+
+    with pytest.raises(
+        TypeError,
+        match="Giving a dimension name to an XTensorVariable indexer is not supported",
+    ):
+        x[("d", idx1),]
+
+
 class TestVectorizedIndexingNotAllowedToBroadcast:
     def test_compile_time_error(self):
         x = xtensor(dims=("a", "b"), shape=(3, 5))