Add xtensor docs

Co-authored-by: Oriol Abril-Pla <[email protected]>

Author: ricardoV94

doc/library/index.rst

@@ -25,6 +25,7 @@ Modules
    sparse/index
    tensor/index
    typed_list
+   xtensor/index
 
 .. module:: pytensor
    :platform: Unix, Windows

doc/library/xtensor/index.md

@@ -0,0 +1,101 @@
+(libdoc_xtensor)=
+# `xtensor` -- XTensor operations
+
+This module implements as abstraction layer on regular tensor operations, that behaves like Xarray.
+
+A new type {class}`pytensor.xtensor.type.XTensorType`, generalizes the {class}`pytensor.tensor.TensorType` 
+with the addition of a `dims` attribute, that labels the dimensions of the tensor. 
+
+Variables of XTensorType (i.e.,  {class}`pytensor.xtensor.type.XTensorVariable`s) are the symbolic counterpart
+to xarray DataArray objects.
+
+The module implements several PyTensor operations {class}`pytensor.xtensor.basic.XOp`s, whose signature mimics that of 
+xarray (and xarray_einstats) DataArray operations. These operations, unlike most regular PyTensor operations, cannot 
+be directly evaluated, but require a rewrite (lowering) into a regular tensor graph that can itself be evaluated as usual.
+
+Like regular PyTensor, we don't need an Op for every possible method or function in the public API of xarray.
+If the existing XOps can be composed to produce the desired result, then we can use them directly.
+
+## Coordinates
+For now, there's no analogous of xarray coordinates, so you won't be able to do coordinate operations like `.sel`.
+The graphs produced by an xarray program without coords are much more amenable to the numpy-like backend of PyTensor.
+Coords involve aspects of Pandas/database query and joining that are not trivially expressible in PyTensor.
+
+## Example
+
+
+```{testcode}
+
+import pytensor.tensor as pt
+import pytensor.xtensor as ptx
+
+a = pt.tensor("a", shape=(3,))
+b = pt.tensor("b", shape=(4,))
+
+ax = ptx.as_xtensor(a, dims=["x"])
+bx = ptx.as_xtensor(b, dims=["y"])
+
+zx = ax + bx
+assert zx.type == ptx.type.XTensorType("float64", dims=["x", "y"], shape=(3, 4))
+
+z = zx.values
+z.dprint()
+```
+
+
+```{testoutput}
+
+TensorFromXTensor [id A]
+ └─ XElemwise{scalar_op=Add()} [id B]
+    ├─ XTensorFromTensor{dims=('x',)} [id C]
+    │  └─ a [id D]
+    └─ XTensorFromTensor{dims=('y',)} [id E]
+       └─ b [id F]
+```
+
+Once we compile the graph, no XOps are left.
+
+```{testcode}
+
+import pytensor
+
+with pytensor.config.change_flags(optimizer_verbose=True):
+    fn = pytensor.function([a, b], z)
+
+```
+
+```{testoutput}
+
+rewriting: rewrite lower_elemwise replaces XElemwise{scalar_op=Add()}.0 of XElemwise{scalar_op=Add()}(XTensorFromTensor{dims=('x',)}.0, XTensorFromTensor{dims=('y',)}.0) with XTensorFromTensor{dims=('x', 'y')}.0 of XTensorFromTensor{dims=('x', 'y')}(Add.0)
+rewriting: rewrite useless_tensor_from_xtensor replaces TensorFromXTensor.0 of TensorFromXTensor(XTensorFromTensor{dims=('x',)}.0) with a of None
+rewriting: rewrite useless_tensor_from_xtensor replaces TensorFromXTensor.0 of TensorFromXTensor(XTensorFromTensor{dims=('y',)}.0) with b of None
+rewriting: rewrite useless_tensor_from_xtensor replaces TensorFromXTensor.0 of TensorFromXTensor(XTensorFromTensor{dims=('x', 'y')}.0) with Add.0 of Add(ExpandDims{axis=1}.0, ExpandDims{axis=0}.0)
+
+```
+
+```{testcode}
+
+fn.dprint()
+```
+
+```{testoutput}
+
+Add [id A] 2
+ ├─ ExpandDims{axis=1} [id B] 1
+ │  └─ a [id C]
+ └─ ExpandDims{axis=0} [id D] 0
+    └─ b [id E]
+```
+
+
+## Index
+
+:::{toctree}
+:maxdepth: 1
+
+module_functions
+math
+linalg
+random
+type
+:::

doc/library/xtensor/linalg.md

@@ -0,0 +1,7 @@
+(libdoc_xtensor_linalg)=
+# `xtensor.linalg` -- Linear algebra operations
+
+```{eval-rst}
+.. automodule:: pytensor.xtensor.linalg
+   :members:
+```

doc/library/xtensor/math.md

@@ -0,0 +1,8 @@
+(libdoc_xtensor_math)=
+# `xtensor.math` Mathematical operations
+
+```{eval-rst}
+.. automodule:: pytensor.xtensor.math
+   :members:
+   :exclude-members: XDot, dot
+```

doc/library/xtensor/module_functions.md

@@ -0,0 +1,7 @@
+(libdoc_xtensor_module_function)=
+# `xtensor` -- Module level operations
+
+```{eval-rst}
+.. automodule:: pytensor.xtensor
+   :members: broadcast, concat, dot, full_like, ones_like, zeros_like
+```

doc/library/xtensor/random.md

@@ -0,0 +1,7 @@
+(libdoc_xtensor_random)=
+# `xtensor.random` Random number generator operations
+
+```{eval-rst}
+.. automodule:: pytensor.xtensor.random
+   :members:
+```

doc/library/xtensor/type.md

@@ -0,0 +1,21 @@
+(libdoc_xtenor_type)=
+
+# `xtensor.type` -- Types and Variables
+
+## XTensorVariable creation functions
+
+```{eval-rst}
+.. automodule:: pytensor.xtensor.type
+   :members: xtensor, xtensor_constant, as_xtensor
+
+```
+
+## XTensor Type and Variable classes
+
+```{eval-rst}
+.. automodule:: pytensor.xtensor.type
+   :noindex:
+   :members: XTensorType, XTensorVariable, XTensorConstant
+```
+
+

pytensor/xtensor/__init__.py

@@ -1,7 +1,7 @@
 import warnings
 
 import pytensor.xtensor.rewriting
-from pytensor.xtensor import linalg, random
+from pytensor.xtensor import linalg, math, random
 from pytensor.xtensor.math import dot
 from pytensor.xtensor.shape import broadcast, concat, full_like, ones_like, zeros_like
 from pytensor.xtensor.type import (

pytensor/xtensor/linalg.py

@@ -11,17 +11,31 @@ def cholesky(
     lower: bool = True,
     *,
     check_finite: bool = False,
-    overwrite_a: bool = False,
     on_error: Literal["raise", "nan"] = "raise",
     dims: Sequence[str],
 ):
+    """Compute the Cholesky decomposition of an XTensorVariable.
+
+    Parameters
+    ----------
+    x : XTensorVariable
+        The input variable to decompose.
+    lower : bool, optional
+        Whether to return the lower triangular matrix. Default is True.
+    check_finite : bool, optional
+        Whether to check that the input is finite. Default is False.
+    on_error : {'raise', 'nan'}, optional
+        What to do if the input is not positive definite. If 'raise', an error is raised.
+        If 'nan', the output will contain NaNs. Default is 'raise'.
+    dims : Sequence[str]
+        The two core dimensions of the input variable, over which the Cholesky decomposition is computed.
+    """
     if len(dims) != 2:
         raise ValueError(f"Cholesky needs two dims, got {len(dims)}")
 
     core_op = Cholesky(
         lower=lower,
         check_finite=check_finite,
-        overwrite_a=overwrite_a,
         on_error=on_error,
     )
     core_dims = (
@@ -40,6 +54,30 @@ def solve(
     lower: bool = False,
     check_finite: bool = False,
 ):
+    """Solve a system of linear equations using XTensorVariables.
+
+    Parameters
+    ----------
+    a : XTensorVariable
+        The left hand-side xtensor.
+    b : XTensorVariable
+        The right-hand side xtensor.
+    dims : Sequence[str]
+        The core dimensions over which to solve the linear equations.
+        If length is 2, we are solving a matrix-vector equation,
+        and the two dimensions should be present in `a`, but only one in `b`.
+        If length is 3, we are solving a matrix-matrix equation,
+        and two dimensions should be present in `a`, two in `b`, and only one should be shared.
+        In both cases the shared dimension will not appear in the output.
+    assume_a : str, optional
+        The type of matrix `a` is assumed to be. Default is 'gen' (general).
+        Options are ["gen", "sym", "her", "pos", "tridiagonal", "banded"].
+        Long form options can also be used ["general", "symmetric", "hermitian", "positive_definite"].
+    lower : bool, optional
+        Whether `a` is lower triangular. Default is False. Only relevant if `assume_a` is "sym", "her", or "pos".
+    check_finite : bool, optional
+        Whether to check that the input is finite. Default is False.
+    """
     a, b = as_xtensor(a), as_xtensor(b)
     input_core_dims: tuple[tuple[str, str], tuple[str] | tuple[str, str]]
     output_core_dims: tuple[tuple[str] | tuple[str, str]]