|
| 1 | +from collections.abc import Sequence |
| 2 | + |
| 3 | +from pytensor.compile import ViewOp |
| 4 | +from pytensor.graph import Apply, Op |
| 5 | +from pytensor.link.c.op import COp |
| 6 | +from pytensor.link.jax.linker import jax_funcify |
| 7 | +from pytensor.link.numba.linker import numba_funcify |
| 8 | +from pytensor.link.pytorch.linker import pytorch_funcify |
| 9 | +from pytensor.tensor.type import TensorType |
| 10 | +from pytensor.xtensor.type import XTensorType, as_xtensor, xtensor |
| 11 | + |
| 12 | + |
| 13 | +class XOp(Op): |
| 14 | + """A base class for XOps that shouldn't be materialized""" |
| 15 | + |
| 16 | + def perform(self, node, inputs, outputs): |
| 17 | + raise NotImplementedError( |
| 18 | + f"xtensor operation {self} must be lowered to equivalent tensor operations" |
| 19 | + ) |
| 20 | + |
| 21 | + |
| 22 | +class XTypeCastOp(COp): |
| 23 | + """Base class for Ops that type cast between TensorType and XTensorType. |
| 24 | +
|
| 25 | + This is like a `ViewOp` but without the expectation the input and output have identical types. |
| 26 | + """ |
| 27 | + |
| 28 | + view_map = {0: [0]} |
| 29 | + |
| 30 | + def perform(self, node, inputs, output_storage): |
| 31 | + output_storage[0][0] = inputs[0] |
| 32 | + |
| 33 | + def c_code(self, node, nodename, inp, out, sub): |
| 34 | + (iname,) = inp |
| 35 | + (oname,) = out |
| 36 | + fail = sub["fail"] |
| 37 | + |
| 38 | + code, _ = ViewOp.c_code_and_version[TensorType] |
| 39 | + return code % locals() |
| 40 | + |
| 41 | + def c_code_cache_version(self): |
| 42 | + _, version = ViewOp.c_code_and_version[TensorType] |
| 43 | + return (version,) |
| 44 | + |
| 45 | + |
| 46 | +@numba_funcify.register(XTypeCastOp) |
| 47 | +def numba_funcify_XCast(op, *args, **kwargs): |
| 48 | + from pytensor.link.numba.dispatch.basic import numba_njit |
| 49 | + |
| 50 | + @numba_njit |
| 51 | + def xcast(x): |
| 52 | + return x |
| 53 | + |
| 54 | + return xcast |
| 55 | + |
| 56 | + |
| 57 | +@jax_funcify.register(XTypeCastOp) |
| 58 | +@pytorch_funcify.register(XTypeCastOp) |
| 59 | +def funcify_XCast(op, *args, **kwargs): |
| 60 | + def xcast(x): |
| 61 | + return x |
| 62 | + |
| 63 | + return xcast |
| 64 | + |
| 65 | + |
| 66 | +class TensorFromXTensor(XTypeCastOp): |
| 67 | + __props__ = () |
| 68 | + |
| 69 | + def make_node(self, x): |
| 70 | + if not isinstance(x.type, XTensorType): |
| 71 | + raise TypeError(f"x must be have an XTensorType, got {type(x.type)}") |
| 72 | + output = TensorType(x.type.dtype, shape=x.type.shape)() |
| 73 | + return Apply(self, [x], [output]) |
| 74 | + |
| 75 | + |
| 76 | +tensor_from_xtensor = TensorFromXTensor() |
| 77 | + |
| 78 | + |
| 79 | +class XTensorFromTensor(XTypeCastOp): |
| 80 | + __props__ = ("dims",) |
| 81 | + |
| 82 | + def __init__(self, dims: Sequence[str]): |
| 83 | + super().__init__() |
| 84 | + self.dims = tuple(dims) |
| 85 | + |
| 86 | + def make_node(self, x): |
| 87 | + if not isinstance(x.type, TensorType): |
| 88 | + raise TypeError(f"x must be an TensorType type, got {type(x.type)}") |
| 89 | + output = xtensor(dtype=x.type.dtype, dims=self.dims, shape=x.type.shape) |
| 90 | + return Apply(self, [x], [output]) |
| 91 | + |
| 92 | + |
| 93 | +def xtensor_from_tensor(x, dims): |
| 94 | + return XTensorFromTensor(dims=dims)(x) |
| 95 | + |
| 96 | + |
| 97 | +class Rename(XTypeCastOp): |
| 98 | + __props__ = ("new_dims",) |
| 99 | + |
| 100 | + def __init__(self, new_dims: tuple[str, ...]): |
| 101 | + super().__init__() |
| 102 | + self.new_dims = new_dims |
| 103 | + |
| 104 | + def make_node(self, x): |
| 105 | + x = as_xtensor(x) |
| 106 | + output = x.type.clone(dims=self.new_dims)() |
| 107 | + return Apply(self, [x], [output]) |
| 108 | + |
| 109 | + |
| 110 | +def rename(x, name_dict: dict[str, str] | None = None, **names: str): |
| 111 | + if name_dict is not None: |
| 112 | + if names: |
| 113 | + raise ValueError("Cannot use both positional and keyword names in rename") |
| 114 | + names = name_dict |
| 115 | + |
| 116 | + x = as_xtensor(x) |
| 117 | + old_names = x.type.dims |
| 118 | + new_names = list(old_names) |
| 119 | + for old_name, new_name in names.items(): |
| 120 | + try: |
| 121 | + new_names[old_names.index(old_name)] = new_name |
| 122 | + except IndexError: |
| 123 | + raise ValueError( |
| 124 | + f"Cannot rename {old_name} to {new_name}: {old_name} not in {old_names}" |
| 125 | + ) |
| 126 | + |
| 127 | + return Rename(tuple(new_names))(x) |
0 commit comments