Implement index operations on XTensorTypes

Author: ricardoV94

pytensor/xtensor/indexing.py

@@ -0,0 +1,97 @@
+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 None (np.newaxis)")
+    idx
+    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_xtensor(idx)
+        except TypeError:
+            idx = as_tensor(idx)
+        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:
+            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 = slice.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 = []
+        for i, idx in enumerate(idxs):
+            if i == x_ndim:
+                raise IndexError("Too many indices")
+            if isinstance(idx.type, XTensorType):
+                raise NotImplementedError(
+                    "Indexing with XTensorType not yet supported."
+                )
+            if isinstance(idx.type, SliceType):
+                out_dims.append(x_dims[i])
+                out_shape.append(get_static_slice_length(idx, x_shape[i]))
+
+            else:  # TensorType
+                if idx.type.ndim == 0:
+                    # Scalar, dimension is dropped
+                    pass
+                elif idx.type.ndim == 1:
+                    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"
+                    )
+
+        output = xtensor(dtype=x.type.dtype, shape=out_shape, dims=out_dims)
+        return Apply(self, [x], [output])
+
+
+index = Index()

pytensor/xtensor/type.py

@@ -1,3 +1,6 @@
+import warnings
+
+
 try:
     import xarray as xr
 
@@ -6,13 +9,13 @@
     XARRAY_AVAILABLE = False
 
 from collections.abc import Sequence
-from typing import TypeVar
+from typing import Any, Literal, TypeVar
 
 import numpy as np
 
 from pytensor import _as_symbolic, config
 from pytensor.graph import Apply, Constant
-from pytensor.graph.basic import Variable, OptionalApplyType
+from pytensor.graph.basic import OptionalApplyType, Variable
 from pytensor.graph.type import HasDataType, HasShape, Type
 from pytensor.tensor.utils import hash_from_ndarray
 from pytensor.utils import hash_from_code
@@ -143,15 +146,50 @@ def __getitem__(self, idx):
 
         return index(self, *idx)
 
+    def sel(self, *args, **kwargs):
+        raise NotImplementedError(
+            "sel not implemented for XTensorVariable, use isel instead"
+        )
 
-class XTensorVariable(Variable):
-    pass
+    def isel(
+        self,
+        indexers: dict[str, Any] | None = None,
+        drop: bool = False,  # Unused by PyTensor
+        missing_dims: Literal["raise", "warn", "ignore"] = "raise",
+        **indexers_kwargs,
+    ):
+        from pytensor.xtensor.indexing import index
+
+        if indexers_kwargs:
+            if indexers is not None:
+                raise ValueError(
+                    "Cannot pass both indexers and indexers_kwargs to isel"
+                )
+            indexers = indexers_kwargs
 
-    # def __str__(self):
-    #     return f"{self.__class__.__name__}{{{self.format},{self.dtype}}}"
+        if missing_dims not in {"raise", "warn", "ignore"}:
+            raise ValueError(
+                f"Unrecognized options {missing_dims} for missing_dims argument"
+            )
 
-    # def __repr__(self):
-    #     return str(self)
+        # Sort indices and pass them to index
+        dims = self.type.dims
+        indices = [slice(None) for _ in self.type.ndim]
+        for key, idx in indexers.items():
+            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 index(self, *indices)
 
 
 class XTensorConstantSignature(tuple):