Extract ViewOp functionality into a base TypeCastOp

Author: ricardoV94

pytensor/compile/ops.py

@@ -33,11 +33,8 @@ def register_view_op_c_code(type, code, version=()):
     ViewOp.c_code_and_version[type] = (code, version)
 
 
-class ViewOp(COp):
-    """
-    Returns an inplace view of the input. Used internally by PyTensor.
-
-    """
+class TypeCastingOp(COp):
+    """Op that performs a graph-level type cast operation, but has no effect computation-wise (identity function)."""
 
     view_map = {0: [0]}
     # Mapping from Type to C code (and version) to use.
@@ -47,13 +44,8 @@ class ViewOp(COp):
     __props__: tuple = ()
     _f16_ok: bool = True
 
-    def make_node(self, x):
-        return Apply(self, [x], [x.type()])
-
-    def perform(self, node, inp, out):
-        (x,) = inp
-        (z,) = out
-        z[0] = x
+    def perform(self, node, inputs, outputs_storage):
+        outputs_storage[0][0] = inputs[0]
 
     def __str__(self):
         return f"{self.__class__.__name__}"
@@ -90,6 +82,13 @@ def c_code_cache_version(self):
 
         return tuple(version)
 
+
+class ViewOp(TypeCastingOp):
+    """Returns an inplace view of the input. Used internally by PyTensor."""
+
+    def make_node(self, x):
+        return Apply(self, [x], [x.type()])
+
     def infer_shape(self, fgraph, node, input_shapes):
         return input_shapes
 

pytensor/link/jax/dispatch/basic.py

@@ -8,7 +8,7 @@
 
 from pytensor.compile import JAX
 from pytensor.compile.builders import OpFromGraph
-from pytensor.compile.ops import DeepCopyOp, ViewOp
+from pytensor.compile.ops import DeepCopyOp, TypeCastingOp
 from pytensor.configdefaults import config
 from pytensor.graph.fg import FunctionGraph
 from pytensor.ifelse import IfElse
@@ -111,12 +111,12 @@ def deepcopyop(x):
     return deepcopyop
 
 
-@jax_funcify.register(ViewOp)
-def jax_funcify_ViewOp(op, **kwargs):
-    def viewop(x):
+@jax_funcify.register(TypeCastingOp)
+def jax_funcify_TypeCastingOp(op, **kwargs):
+    def type_cast(x):
         return x
 
-    return viewop
+    return type_cast
 
 
 @jax_funcify.register(OpFromGraph)

pytensor/link/numba/dispatch/scalar.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 
-from pytensor.compile.ops import ViewOp
+from pytensor.compile.ops import TypeCastingOp
 from pytensor.graph.basic import Variable
 from pytensor.link.numba.dispatch import basic as numba_basic
 from pytensor.link.numba.dispatch.basic import (
@@ -198,14 +198,14 @@ def cast(x):
 
 
 @numba_basic.numba_njit
-def viewop(x):
+def identity(x):
     return x
 
 
 @numba_funcify.register(Identity)
-@numba_funcify.register(ViewOp)
-def numba_funcify_ViewOp(op, **kwargs):
-    return numba_basic.global_numba_func(viewop)
+@numba_funcify.register(TypeCastingOp)
+def numba_funcify_type_casting(op, **kwargs):
+    return numba_basic.global_numba_func(identity)
 
 
 @numba_basic.numba_njit

pytensor/link/pytorch/dispatch/basic.py

@@ -9,7 +9,7 @@
 from pytensor.compile import PYTORCH
 from pytensor.compile.builders import OpFromGraph
 from pytensor.compile.function.types import add_supervisor_to_fgraph
-from pytensor.compile.ops import DeepCopyOp
+from pytensor.compile.ops import DeepCopyOp, TypeCastingOp
 from pytensor.graph.basic import Constant
 from pytensor.graph.fg import FunctionGraph
 from pytensor.ifelse import IfElse
@@ -71,6 +71,14 @@ def pytorch_funcify_FunctionGraph(
     )
 
 
+@pytorch_funcify.register(TypeCastingOp)
+def pytorch_funcify_CastingOp(op, node, **kwargs):
+    def type_cast(x):
+        return x
+
+    return type_cast
+
+
 @pytorch_funcify.register(CheckAndRaise)
 def pytorch_funcify_CheckAndRaise(op, **kwargs):
     error = op.exc_type