replace pytensorf identity with pytensor identity

Author: ferrine

pymc/pytensorf.py

@@ -22,7 +22,6 @@
 import pytensor.tensor as pt
 import scipy.sparse as sps
 
-from pytensor import scalar
 from pytensor.compile import Function, Mode, get_mode
 from pytensor.compile.builders import OpFromGraph
 from pytensor.gradient import grad
@@ -38,6 +37,7 @@
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.op import Op
 from pytensor.scalar.basic import Cast
+from pytensor.scalar.basic import identity as scalar_identity
 from pytensor.scan.op import Scan
 from pytensor.tensor.basic import _as_tensor_variable
 from pytensor.tensor.elemwise import Elemwise
@@ -412,28 +412,6 @@ def hessian_diag(f, vars=None, negate_output=True):
         return empty_gradient
 
 
-class IdentityOp(scalar.UnaryScalarOp):
-    @staticmethod
-    def st_impl(x):
-        return x
-
-    def impl(self, x):
-        return x
-
-    def grad(self, inp, grads):
-        return grads
-
-    def c_code(self, node, name, inp, out, sub):
-        return f"{out[0]} = {inp[0]};"
-
-    def __eq__(self, other):
-        return isinstance(self, type(other))
-
-    def __hash__(self):
-        return hash(type(self))
-
-
-scalar_identity = IdentityOp(scalar.upgrade_to_float, name="scalar_identity")
 identity = Elemwise(scalar_identity, name="identity")
 
 

pymc/sampling/jax.py

@@ -46,7 +46,6 @@
 from pymc.distributions.multivariate import PosDefMatrix
 from pymc.initial_point import StartDict
 from pymc.logprob.utils import CheckParameterValue
-from pymc.pytensorf import IdentityOp
 from pymc.sampling.mcmc import _init_jitter
 from pymc.stats.convergence import log_warnings, run_convergence_checks
 from pymc.util import (
@@ -70,14 +69,6 @@
 )
 
 
-@jax_funcify.register(IdentityOp)
-def jax_funcify_Identity(op, **kwargs):
-    def identity_fn(value):
-        return value
-
-    return identity_fn
-
-
 @jax_funcify.register(Assert)
 @jax_funcify.register(CheckParameterValue)
 def jax_funcify_Assert(op, **kwargs):