WIP Implement index operations for XTensorVariables

Author: ricardoV94

pytensor/xtensor/__init__.py

@@ -7,7 +7,6 @@
 )
 from pytensor.xtensor.shape import concat
 from pytensor.xtensor.type import (
-    XTensorType,
     as_xtensor,
     xtensor,
     xtensor_constant,

pytensor/xtensor/indexing.py

@@ -0,0 +1,142 @@
+# HERE LIE DRAGONS
+# Uselful links to make sense of all the numpy/xarray complexity
+# https://numpy.org/devdocs//user/basics.indexing.html
+# https://numpy.org/neps/nep-0021-advanced-indexing.html
+# https://docs.xarray.dev/en/latest/user-guide/indexing.html
+# https://tutorial.xarray.dev/intermediate/indexing/advanced-indexing.html
+
+from pytensor.graph.basic import Apply, Constant, Variable
+from pytensor.scalar.basic import discrete_dtypes
+from pytensor.tensor.basic import as_tensor
+from pytensor.tensor.type_other import NoneTypeT, SliceType, make_slice
+from pytensor.xtensor.basic import XOp
+from pytensor.xtensor.type import XTensorType, as_xtensor, xtensor
+
+
+def as_idx_variable(idx):
+    if idx is None or (isinstance(idx, Variable) and isinstance(idx.type, NoneTypeT)):
+        raise TypeError(
+            "XTensors do not support indexing with None (np.newaxis), use expand_dims instead"
+        )
+    if isinstance(idx, slice):
+        idx = make_slice(idx)
+    elif isinstance(idx, Variable) and isinstance(idx.type, SliceType):
+        pass
+    else:
+        # Must be integer indices, we already counted for None and slices
+        try:
+            idx = as_tensor(idx)
+        except TypeError:
+            idx = as_xtensor(idx)
+        if idx.type.dtype == "bool":
+            raise NotImplementedError("Boolean indexing not yet supported")
+        if idx.type.dtype not in discrete_dtypes:
+            raise TypeError("Numerical indices must be integers or boolean")
+        if idx.type.dtype == "bool" and idx.type.ndim == 0:
+            # This can't be triggered right now, but will once we lift the boolean restriction
+            raise NotImplementedError("Scalar boolean indices not supported")
+    return idx
+
+
+def get_static_slice_length(slc: Variable, dim_length: None | int) -> int | None:
+    if dim_length is None:
+        return None
+    if isinstance(slc, Constant):
+        d = slc.data
+        start, stop, step = d.start, d.stop, d.step
+    elif slc.owner is None:
+        # It's a root variable no way of knowing what we're getting
+        return None
+    else:
+        # It's a MakeSliceOp
+        start, stop, step = slc.owner.inputs
+        if isinstance(start, Constant):
+            start = start.data
+        else:
+            return None
+        if isinstance(stop, Constant):
+            stop = stop.data
+        else:
+            return None
+        if isinstance(step, Constant):
+            step = step.data
+        else:
+            return None
+    return len(range(*slice(start, stop, step).indices(dim_length)))
+
+
+class Index(XOp):
+    __props__ = ()
+
+    def make_node(self, x, *idxs):
+        x = as_xtensor(x)
+        idxs = [as_idx_variable(idx) for idx in idxs]
+
+        x_ndim = x.type.ndim
+        x_dims = x.type.dims
+        x_shape = x.type.shape
+        out_dims = []
+        out_shape = []
+        has_unlabeled_vector_idx = False
+        has_labeled_vector_idx = False
+        for i, idx in enumerate(idxs):
+            if i == x_ndim:
+                raise IndexError("Too many indices")
+            if isinstance(idx.type, SliceType):
+                out_dims.append(x_dims[i])
+                out_shape.append(get_static_slice_length(idx, x_shape[i]))
+            elif isinstance(idx.type, XTensorType):
+                if has_unlabeled_vector_idx:
+                    raise NotImplementedError(
+                        "Mixing of labeled and unlabeled vector indexing not implemented"
+                    )
+                has_labeled_vector_idx = True
+                idx_dims = idx.type.dims
+                for dim in idx_dims:
+                    idx_dim_shape = idx.type.shape[idx_dims.index(dim)]
+                    if dim in out_dims:
+                        # Dim already introduced in output by a previous index
+                        # Update static shape or raise if incompatible
+                        out_dim_pos = out_dims.index(dim)
+                        out_dim_shape = out_shape[out_dim_pos]
+                        if out_dim_shape is None:
+                            # We don't know the size of the dimension yet
+                            out_shape[out_dim_pos] = idx_dim_shape
+                        elif (
+                            idx_dim_shape is not None and idx_dim_shape != out_dim_shape
+                        ):
+                            raise IndexError(
+                                f"Dimension of indexers mismatch for dim {dim}"
+                            )
+                    else:
+                        # New dimension
+                        out_dims.append(dim)
+                        out_shape.append(idx_dim_shape)
+
+            else:  # TensorType
+                if idx.type.ndim == 0:
+                    # Scalar, dimension is dropped
+                    pass
+                elif idx.type.ndim == 1:
+                    if has_labeled_vector_idx:
+                        raise NotImplementedError(
+                            "Mixing of labeled and unlabeled vector indexing not implemented"
+                        )
+                    has_unlabeled_vector_idx = True
+                    out_dims.append(x_dims[i])
+                    out_shape.append(idx.type.shape[0])
+                else:
+                    # Same error that xarray raises
+                    raise IndexError(
+                        "Unlabeled multi-dimensional array cannot be used for indexing"
+                    )
+        for j in range(i + 1, x_ndim):
+            # Add any unindexed dimensions
+            out_dims.append(x_dims[j])
+            out_shape.append(x_shape[j])
+
+        output = xtensor(dtype=x.type.dtype, shape=out_shape, dims=out_dims)
+        return Apply(self, [x, *idxs], [output])
+
+
+index = Index()

pytensor/xtensor/rewriting/__init__.py

@@ -1,4 +1,5 @@
 import pytensor.xtensor.rewriting.basic
+import pytensor.xtensor.rewriting.indexing
 import pytensor.xtensor.rewriting.reduction
 import pytensor.xtensor.rewriting.shape
 import pytensor.xtensor.rewriting.vectorization

pytensor/xtensor/rewriting/indexing.py

@@ -0,0 +1,67 @@
+from pytensor.graph import Constant, node_rewriter
+from pytensor.tensor import TensorType, specify_shape
+from pytensor.tensor.type_other import NoneTypeT, SliceType
+from pytensor.xtensor.basic import tensor_from_xtensor, xtensor_from_tensor
+from pytensor.xtensor.indexing import Index
+from pytensor.xtensor.rewriting.utils import register_xcanonicalize
+from pytensor.xtensor.type import XTensorType
+
+
+def to_basic_idx(idx):
+    if isinstance(idx.type, SliceType):
+        if isinstance(idx, Constant):
+            return idx.data
+        elif idx.owner:
+            # MakeSlice Op
+            # We transform NoneConsts to regular None so that basic Subtensor can be used if possible
+            return slice(
+                *[
+                    None if isinstance(i.type, NoneTypeT) else i
+                    for i in idx.owner.inputs
+                ]
+            )
+        else:
+            return idx
+    if (
+        isinstance(idx.type, XTensorType | TensorType)
+        and idx.type.ndim == 0
+        and idx.type.dtype != bool
+    ):
+        return idx
+    raise TypeError("Cannot convert idx to basic idx")
+
+
+def _count_idx_types(idxs):
+    basic, vector, xvector = 0, 0, 0
+    for idx in idxs:
+        if isinstance(idx.type, SliceType):
+            basic += 1
+        elif idx.type.ndim == 0:
+            basic += 1
+        elif isinstance(idx.type, TensorType):
+            vector += 1
+        else:
+            xvector += 1
+    return basic, vector, xvector
+
+
+@register_xcanonicalize
+@node_rewriter(tracks=[Index])
+def lower_index(fgraph, node):
+    x, *idxs = node.inputs
+    [out] = node.outputs
+    x_tensor = tensor_from_xtensor(x)
+    n_basic, n_vector, n_xvector = _count_idx_types(idxs)
+    if n_xvector == 0 and n_vector == 0:
+        x_tensor_indexed = x_tensor[tuple(to_basic_idx(idx) for idx in idxs)]
+    elif n_vector == 1 and n_xvector == 0:
+        # Special case for single vector index, no orthogonal indexing
+        x_tensor_indexed = x_tensor[tuple(idxs)]
+    else:
+        # Not yet implemented
+        return None
+
+    # Add lost shape if any
+    x_tensor_indexed = specify_shape(x_tensor_indexed, out.type.shape)
+    new_out = xtensor_from_tensor(x_tensor_indexed, dims=out.type.dims)
+    return [new_out]

pytensor/xtensor/type.py

@@ -1,3 +1,5 @@
+import warnings
+
 from pytensor.tensor import TensorType
 from pytensor.tensor.math import variadic_mul
 
@@ -10,7 +12,7 @@
     XARRAY_AVAILABLE = False
 
 from collections.abc import Sequence
-from typing import TypeVar
+from typing import Any, Literal, TypeVar
 
 import numpy as np
 
@@ -339,7 +341,112 @@ def sel(self, *args, **kwargs):
         raise NotImplementedError("sel not implemented for XTensorVariable")
 
     def __getitem__(self, idx):
-        raise NotImplementedError("Indexing not yet implemnented")
+        if isinstance(idx, dict):
+            return self.isel(idx)
+
+        # Check for ellipsis not in the last position (last one is useless anyway)
+        if any(idx_item is Ellipsis for idx_item in idx):
+            if idx.count(Ellipsis) > 1:
+                raise IndexError("an index can only have a single ellipsis ('...')")
+            # Convert intermediate Ellipsis to slice(None)
+            ellipsis_loc = idx.index(Ellipsis)
+            n_implied_none_slices = self.type.ndim - (len(idx) - 1)
+            idx = (
+                *idx[:ellipsis_loc],
+                *((slice(None),) * n_implied_none_slices),
+                *idx[ellipsis_loc + 1 :],
+            )
+
+        return px.indexing.index(self, *idx)
+
+    def isel(
+        self,
+        indexers: dict[str, Any] | None = None,
+        drop: bool = False,  # Unused by PyTensor
+        missing_dims: Literal["raise", "warn", "ignore"] = "raise",
+        **indexers_kwargs,
+    ):
+        if indexers_kwargs:
+            if indexers is not None:
+                raise ValueError(
+                    "Cannot pass both indexers and indexers_kwargs to isel"
+                )
+            indexers = indexers_kwargs
+
+        if missing_dims not in {"raise", "warn", "ignore"}:
+            raise ValueError(
+                f"Unrecognized options {missing_dims} for missing_dims argument"
+            )
+
+        # Sort indices and pass them to index
+        dims = self.type.dims
+        indices = [slice(None)] * self.type.ndim
+        for key, idx in indexers.items():
+            if idx is Ellipsis:
+                # Xarray raises a less informative error, suggesting indices must be integer
+                # But slices are also fine
+                raise TypeError("Ellipsis (...) is an invalid labeled index")
+            try:
+                indices[dims.index(key)] = idx
+            except IndexError:
+                if missing_dims == "raise":
+                    raise ValueError(
+                        f"Dimension {key} does not exist. Expected one of {dims}"
+                    )
+                elif missing_dims == "warn":
+                    warnings.warn(
+                        UserWarning,
+                        f"Dimension {key} does not exist. Expected one of {dims}",
+                    )
+
+        return px.indexing.index(self, *indices)
+
+    def _head_tail_or_thin(
+        self,
+        indexers: dict[str, Any] | int | None,
+        indexers_kwargs: dict[str, Any],
+        *,
+        kind: Literal["head", "tail", "thin"],
+    ):
+        if indexers_kwargs:
+            if indexers is not None:
+                raise ValueError(
+                    "Cannot pass both indexers and indexers_kwargs to head"
+                )
+            indexers = indexers_kwargs
+
+        if indexers is None:
+            if kind == "thin":
+                raise TypeError(
+                    "thin() indexers must be either dict-like or a single integer"
+                )
+            else:
+                # Default to 5 for head and tail
+                indexers = {dim: 5 for dim in self.type.dims}
+
+        elif not isinstance(indexers, dict):
+            indexers = {dim: indexers for dim in self.type.dims}
+
+        if kind == "head":
+            indices = {dim: slice(None, value) for dim, value in indexers.items()}
+        elif kind == "tail":
+            sizes = self.sizes
+            # Can't use slice(-value, None), in case value is zero
+            indices = {
+                dim: slice(sizes[dim] - value, None) for dim, value in indexers.items()
+            }
+        elif kind == "thin":
+            indices = {dim: slice(None, None, value) for dim, value in indexers.items()}
+        return self.isel(indices)
+
+    def head(self, indexers: dict[str, Any] | int | None = None, **indexers_kwargs):
+        return self._head_tail_or_thin(indexers, indexers_kwargs, kind="head")
+
+    def tail(self, indexers: dict[str, Any] | int | None = None, **indexers_kwargs):
+        return self._head_tail_or_thin(indexers, indexers_kwargs, kind="tail")
+
+    def thin(self, indexers: dict[str, Any] | int | None = None, **indexers_kwargs):
+        return self._head_tail_or_thin(indexers, indexers_kwargs, kind="thin")
 
     # ndarray methods
     # https://docs.xarray.dev/en/latest/api.html#id7

tests/xtensor/test_indexing.py

@@ -0,0 +1,42 @@
+import numpy as np
+import pytest
+from xarray import DataArray
+
+from pytensor.xtensor import xtensor
+from tests.xtensor.util import xr_assert_allclose, xr_function
+
+
+@pytest.mark.parametrize(
+    "indices",
+    [
+        (0,),
+        (slice(1, None),),
+        (slice(None, -1),),
+        (slice(None, None, -1),),
+        (0, slice(None), -1, slice(1, None)),
+        (..., 0, -1),
+        (0, ..., -1),
+        (0, -1, ...),
+    ],
+)
+@pytest.mark.parametrize("labeled", (False, True), ids=["unlabeled", "labeled"])
+def test_basic_indexing(labeled, indices):
+    if ... in indices and labeled:
+        pytest.skip("Ellipsis not supported with labeled indexing")
+
+    dims = ("a", "b", "c", "d")
+    x = xtensor(dims=dims, shape=(2, 3, 5, 7))
+
+    if labeled:
+        shufled_dims = tuple(np.random.permutation(dims))
+        indices = dict(zip(shufled_dims, indices, strict=False))
+    out = x[indices]
+
+    fn = xr_function([x], out)
+    x_test_values = np.arange(np.prod(x.type.shape), dtype=x.type.dtype).reshape(
+        x.type.shape
+    )
+    x_test = DataArray(x_test_values, dims=x.type.dims)
+    res = fn(x_test)
+    expected_res = x_test[indices]
+    xr_assert_allclose(res, expected_res)