work on tests

Author: aseyboldt

pytensor/xtensor/basic.py

@@ -5,7 +5,7 @@
 from pytensor.tensor.basic import zeros as tensor_zeros
 from pytensor.tensor.shape import specify_shape
 from pytensor.tensor.type import TensorType
-from pytensor.xtensor.type import DimVariable, XTensorType, as_xtensor, xtensor
+from pytensor.xtensor.type import DimVariable, XTensorType, as_dim, as_xtensor, xtensor
 
 
 DIM_LENGTH_SCALAR = uint64
@@ -64,6 +64,7 @@ def L_op(self, inputs, outs, g_outs):
 def xtensor_from_tensor(x, dims, name=None, check: bool = True):
     if check:
         x = specify_shape(x, [dim.size for dim in dims])
+    dims = [as_dim(dim) for dim in dims]
     return XTensorFromTensor()(x, *dims, name=name)
 
 

pytensor/xtensor/dims.py

@@ -41,6 +41,8 @@ def perform(self, node, inputs, outputs):
 
 
 def _dim_size(dim: DimVariable) -> DIM_LENGTH_VARIABLE:
+    if dim.type.size is not None:
+        return DIM_LENGTH_TYPE.filter_variable(dim.type.size)
     return Length()(dim)
 
 
@@ -162,21 +164,26 @@ def product_dim(*dims: DimVariable, name: str | None = None) -> DimVariable:
     return Product()(*dims, name=name)
 
 
-def rebase_dim(dim: DimVariable, *tensors: XTensorVariable) -> DimVariable:
-    if not isinstance(dim, DimVariable):
+def rebase_dim(dim: DimVariable | DimType, *tensors: XTensorVariable) -> DimVariable:
+    if not isinstance(dim, DimVariable | DimType):
         raise TypeError(f"dim must be a DimVariable, got {type(dim)}")
 
     if not tensors:
         raise ValueError("At least one tensor must be provided for rebasing.")
 
+    if isinstance(dim, DimVariable):
+        dim_type = dim.type
+    else:
+        dim_type = dim
+
     for tensor in tensors:
         for i, tensor_dim in enumerate(tensor.type.dims):
-            if dim.type == tensor_dim:
+            if dim_type == tensor_dim:
                 return _dim_from_tensor(tensor, idx=i)
-    raise ValueError(f"Dimension {dim.type} not found in any of the provided tensors.")
+    raise ValueError(f"Dimension {dim} not found in any of the provided tensors.")
 
 
 def rebase_dims(
-    dims: Iterable[DimVariable], *tensors: XTensorVariable
+    dims: Iterable[DimVariable | DimType], *tensors: XTensorVariable
 ) -> list[DimVariable]:
     return [rebase_dim(dim, *tensors) for dim in dims]

pytensor/xtensor/math.py

@@ -9,7 +9,8 @@
 from pytensor.graph.basic import Apply
 from pytensor.scalar.basic import _cast_mapping, upcast
 from pytensor.xtensor.basic import XOp, as_xtensor
-from pytensor.xtensor.type import AsDim, DimType, DimVariable, as_dim, xtensor
+from pytensor.xtensor.dims import rebase_dims
+from pytensor.xtensor.type import AsDim, DimType, as_dim_type, xtensor
 from pytensor.xtensor.vectorization import XElemwise
 
 
@@ -527,7 +528,7 @@ class Dot(XOp):
     __props__ = ("dims",)
 
     def __init__(self, dims: Iterable[DimType]):
-        self.dims = dims
+        self.dims = frozenset(dims)
         super().__init__()
 
     def make_node(self, x, y):
@@ -544,7 +545,7 @@ def make_node(self, x, y):
         # Determine output dtype
         out_dtype = upcast(x.type.dtype, y.type.dtype)
 
-        out = xtensor(dtype=out_dtype, dims=out_dims)
+        out = xtensor(dtype=out_dtype, dims=rebase_dims(out_dims, x, y))
         return Apply(self, [x, y], [out])
 
 
@@ -589,7 +590,6 @@ def dot(x, y, dim: str | Iterable[AsDim] | EllipsisType | None = None):
 
     x_dims = set(x.type.dims)
     y_dims = set(y.type.dims)
-    print(x_dims, y_dims, dim)
     intersection = x_dims & y_dims
     union = x_dims | y_dims
 
@@ -599,9 +599,9 @@ def dot(x, y, dim: str | Iterable[AsDim] | EllipsisType | None = None):
     elif dim is ...:
         dim_set = union
     elif isinstance(dim, Iterable):
-        dim_set = set([as_dim(dim).type for dim in dim])
-    elif isinstance(dim, (DimVariable, DimType, str)):
-        dim_set = {as_dim(dim).type}
+        dim_set = {as_dim_type(dim) for dim in dim}
+    elif isinstance(dim, AsDim):
+        dim_set = {as_dim_type(dim)}
     else:
         raise TypeError(f"Unknown type {dim} for dimension")
 

pytensor/xtensor/reduction.py

@@ -7,12 +7,20 @@
 from pytensor.graph.basic import Apply
 from pytensor.tensor.math import variadic_mul
 from pytensor.xtensor.basic import XOp
+from pytensor.xtensor.dims import rebase_dims
 from pytensor.xtensor.math import neq, sqrt
 from pytensor.xtensor.math import sqr as square
-from pytensor.xtensor.type import DimType, DimVariable, as_xtensor, xtensor
+from pytensor.xtensor.type import (
+    AsDim,
+    DimType,
+    DimVariable,
+    as_dim_type,
+    as_xtensor,
+    xtensor,
+)
 
 
-REDUCE_DIM = DimVariable | Sequence[DimVariable] | EllipsisType | None
+REDUCE_DIM = DimVariable | Sequence[AsDim] | EllipsisType | None
 
 
 class XReduce(XOp):
@@ -22,28 +30,28 @@ def __init__(self, binary_op, dims: Sequence[DimVariable]):
         super().__init__()
         self.binary_op = binary_op
         # Order of reduce dims doesn't change the behavior of the Op
-        self.dims = tuple(dims)
+        self.dims = frozenset(dims)
 
     def make_node(self, x):
         x = as_xtensor(x)
         x_dims = x.type.dims
         x_dims_set = set(x_dims)
         reduce_dims_set = set(self.dims)
         if x_dims_set == reduce_dims_set:
-            out_dims, out_shape = [], []
+            out_dim_types, out_shape = [], []
         else:
             if not reduce_dims_set.issubset(x_dims_set):
                 raise ValueError(
                     f"Reduced dims {self.dims} not found in array dimensions {x_dims}."
                 )
-            out_dims, out_shape = zip(
+            out_dim_types, out_shape = zip(
                 *[
                     (d, s)
                     for d, s in zip(x_dims, x.type.shape)
                     if d not in reduce_dims_set
                 ]
             )
-        output = xtensor(dtype=x.type.dtype, dims=out_dims)
+        output = xtensor(dtype=x.type.dtype, dims=rebase_dims(out_dim_types, x))
         return Apply(self, [x], [output])
 
 
@@ -53,7 +61,7 @@ def _process_user_dims(x, dim: REDUCE_DIM) -> Sequence[DimType]:
     elif dim is None or dim is Ellipsis:
         x = as_xtensor(x)
         return typing.cast(tuple[DimType], x.type.dims)
-    return tuple(dim.type for dim in dim)
+    return tuple(as_dim_type(dim) for dim in dim)
 
 
 def reduce(x, dim: REDUCE_DIM = None, *, binary_op):
@@ -112,9 +120,9 @@ def std(x, dim: REDUCE_DIM, *, ddof: int = 0):
 class XCumReduce(XOp):
     __props__ = ("binary_op", "dims")
 
-    def __init__(self, binary_op, dims: Sequence[str]):
+    def __init__(self, binary_op, dims: Sequence[DimType]):
         self.binary_op = binary_op
-        self.dims = tuple(sorted(dims))  # Order doesn't matter
+        self.dims = frozenset(dims)
 
     def make_node(self, x):
         x = as_xtensor(x)

pytensor/xtensor/rewriting/basic.py

@@ -51,6 +51,18 @@ def useless_length(fgraph, node):
         return [dim.owner.inputs[0]]
 
 
+@register_infer_shape
+@register_useless
+@register_canonicalize
+@register_lower_xtensor
+@node_rewriter(tracks=[Length])
+def known_length(fgraph, node):
+    """Length(dim_with_size) -> size"""
+    [dim] = node.inputs
+    if dim.type.size is not None:
+        return [dim.type.size]
+
+
 @register_infer_shape
 @register_useless
 @register_canonicalize

pytensor/xtensor/rewriting/reduction.py

@@ -17,6 +17,7 @@ def lower_reduce(fgraph, node):
     x_dims = x.type.dims
     reduce_dims = node.op.dims
     reduce_axis = [x_dims.index(dim) for dim in reduce_dims]
+    out_dims = [x_dim for x_dim in x.dims if x_dim.type not in reduce_dims]
 
     if not reduce_axis:
         return [x]
@@ -40,7 +41,7 @@ def lower_reduce(fgraph, node):
 
     x_tensor = tensor_from_xtensor(x)
     out_tensor = tensor_op_class(axis=reduce_axis)(x_tensor)
-    new_out = xtensor_from_tensor(out_tensor, out.type.dims)
+    new_out = xtensor_from_tensor(out_tensor, out_dims)
     return [new_out]
 
 
@@ -51,6 +52,7 @@ def lower_cumreduce(fgraph, node):
     x_dims = x.type.dims
     reduce_dims = node.op.dims
     reduce_axis = [x_dims.index(dim) for dim in reduce_dims]
+    out_dims = [x_dim for x_dim in x.dims if x_dim not in reduce_dims]
 
     if not reduce_axis:
         return [x]
@@ -68,5 +70,5 @@ def lower_cumreduce(fgraph, node):
     out_tensor = tensor_from_xtensor(x)
     for axis in reduce_axis:
         out_tensor = tensor_op_class(axis=axis)(out_tensor)
-    out = xtensor_from_tensor(out_tensor, x.type.dims)
+    out = xtensor_from_tensor(out_tensor, out_dims)
     return [out]

pytensor/xtensor/rewriting/utils.py

@@ -6,7 +6,7 @@
 from pytensor.graph.rewriting.db import EquilibriumDB, RewriteDatabase
 from pytensor.tensor.rewriting.ofg import inline_ofg_expansion
 from pytensor.tensor.variable import TensorVariable
-from pytensor.xtensor.type import XTensorVariable
+from pytensor.xtensor.type import AsDim, XTensorVariable, as_dim_type
 
 
 lower_xtensor_db = EquilibriumDB(ignore_newtrees=False)
@@ -56,8 +56,9 @@ def register(inner_rewriter: RewriteDatabase | NodeRewriter):
         return node_rewriter
 
 
-def lower_aligned(x: XTensorVariable, out_dims: Sequence[str]) -> TensorVariable:
+def lower_aligned(x: XTensorVariable, out_dims: Sequence[AsDim]) -> TensorVariable:
     """Lower an XTensorVariable to a TensorVariable so that it's dimensions are aligned with "out_dims"."""
+    out_dim_types = [as_dim_type(x) for x in out_dims]
     inp_dims = {d: i for i, d in enumerate(x.type.dims)}
-    ds_order = tuple(inp_dims.get(dim, "x") for dim in out_dims)
+    ds_order = tuple(inp_dims.get(dim, "x") for dim in out_dim_types)
     return typing.cast(TensorVariable, x.values.dimshuffle(ds_order))

pytensor/xtensor/shape.py

@@ -22,7 +22,6 @@
     as_xtensor,
     xtensor,
 )
-from pytensor.xtensor.vectorization import combine_dims_and_shape
 
 
 class Stack(XOp):

pytensor/xtensor/type.py

@@ -1138,16 +1138,36 @@ def as_symbolic_xarray(x, **kwargs):
 
 def as_dim(
     x: AsDim,
+    *,
+    allow_new: bool = False,
 ) -> DimVariable:
     if isinstance(x, DimVariable):
         return x
     if isinstance(x, str):
-        return dim(name=x, unique=False)
+        if allow_new:
+            return dim(name=x, unique=True)
+        else:
+            raise ValueError(
+                f"Cannot convert string {x} to dim without allow_new=True. "
+                "Use `dim(name=x)` to create a new dimension."
+            )
     if isinstance(x, DimType):
-        return cast(DimVariable, x())
+        if allow_new:
+            return cast(DimVariable, x())
+        else:
+            raise ValueError(
+                f"Cannot convert DimType {x} to dim without allow_new=True. "
+                "Use `x.make_variable()` to create a new dimension variable."
+            )
     raise ValueError(f"Can not convert {type(x)} to dim.")
 
 
+def as_dim_type(
+    x: AsDim,
+) -> DimType:
+    return as_dim(x, allow_new=True).type
+
+
 def as_xtensor(
     x,
     dims: Sequence[DimVariable] | None = None,

pytensor/xtensor/vectorization.py

@@ -19,6 +19,7 @@
     DimType,
     DimVariable,
     XTensorVariable,
+    as_dim,
     as_xtensor,
     xtensor,
 )
@@ -27,12 +28,15 @@
 def broadcast_xtensors(
     inputs: Sequence[XTensorVariable], exclude: Sequence[AsDim] | None = None
 ) -> list[DimVariable]:
+    if exclude is None:
+        exclude = []
+    exclude_set: set[DimType] = {as_dim(d).type for d in exclude}
     dims_and_shape: dict[DimType, int | None] = {}
     dim_to_dimvar: dict[DimType, DimVariable] = {}
     for inp in inputs:
         for dim, dim_length in zip(inp.dims, inp.type.shape):
             # TODO Must check dim conversion!!!
-            if dim in exclude:
+            if dim.type in exclude_set:
                 continue
             if dim.type not in dims_and_shape:
                 dims_and_shape[dim.type] = dim_length