Implement casting for XTensorVariables

Author: ricardoV94

pytensor/xtensor/math.py

@@ -1,8 +1,13 @@
 import inspect
 import sys
 
+import numpy as np
+
 import pytensor.scalar as ps
+from pytensor import config
 from pytensor.scalar import ScalarOp
+from pytensor.scalar.basic import _cast_mapping
+from pytensor.xtensor.basic import as_xtensor
 from pytensor.xtensor.vectorization import XElemwise
 
 
@@ -29,3 +34,26 @@ def get_all_scalar_ops():
 
 for name, op in get_all_scalar_ops().items():
     setattr(this_module, name, op)
+
+
+_xelemwise_cast_op: dict[str, XElemwise] = {}
+
+
+def cast(x, dtype):
+    if dtype == "floatX":
+        dtype = config.floatX
+    else:
+        dtype = np.dtype(dtype).name
+
+    x = as_xtensor(x)
+    if x.type.dtype == dtype:
+        return x
+    if x.type.dtype.startswith("complex") and not dtype.startswith("complex"):
+        raise TypeError(
+            "Casting from complex to real is ambiguous: consider"
+            " real(), imag(), angle() or abs()"
+        )
+
+    if dtype not in _xelemwise_cast_op:
+        _xelemwise_cast_op[dtype] = XElemwise(scalar_op=_cast_mapping[dtype])
+    return _xelemwise_cast_op[dtype](x)

tests/xtensor/test_math.py

@@ -1,5 +1,6 @@
 # ruff: noqa: E402
 import pytest
+from xtensor.util import xr_arange_like
 
 
 pytest.importorskip("xarray")  #
@@ -107,3 +108,21 @@ def test_multiple_constant():
     res = fn(x_test)
     expected_res = np.exp(x_test * 2) + 2
     np.testing.assert_allclose(res, expected_res)
+
+
+def test_cast():
+    x = xtensor("x", shape=(2, 3), dims=("a", "b"), dtype="float32")
+    yf64 = x.astype("float64")
+    yi16 = x.astype("int16")
+    ybool = x.astype("bool")
+
+    fn = xr_function([x], [yf64, yi16, ybool])
+    x_test = xr_arange_like(x)
+    res_f64, res_i16, res_bool = fn(x_test)
+    xr_assert_allclose(res_f64, x_test.astype("float64"))
+    xr_assert_allclose(res_i16, x_test.astype("int16"))
+    xr_assert_allclose(res_bool, x_test.astype("bool"))
+
+    yc64 = x.astype("complex64")
+    with pytest.raises(TypeError, match="Casting from complex to real is ambiguous"):
+        yc64.astype("float64")