Rename remaining instances of aes and aer

Author: jessegrabowski

pytensor/compile/profiling.py

@@ -1485,65 +1485,65 @@ def print_tips(self, file):
             file=file,
         )
 
-        from pytensor import scalar as aes
+        from pytensor import scalar as ps
         from pytensor.tensor.elemwise import Elemwise
         from pytensor.tensor.math import Dot
 
         scalar_op_amdlibm_no_speed_up = [
-            aes.LT,
-            aes.GT,
-            aes.LE,
-            aes.GE,
-            aes.EQ,
-            aes.NEQ,
-            aes.InRange,
-            aes.Switch,
-            aes.OR,
-            aes.XOR,
-            aes.AND,
-            aes.Invert,
-            aes.ScalarMaximum,
-            aes.ScalarMinimum,
-            aes.Add,
-            aes.Mul,
-            aes.Sub,
-            aes.TrueDiv,
-            aes.IntDiv,
-            aes.Clip,
-            aes.Second,
-            aes.Identity,
-            aes.Cast,
-            aes.Sign,
-            aes.Neg,
-            aes.Reciprocal,
-            aes.Sqr,
+            ps.LT,
+            ps.GT,
+            ps.LE,
+            ps.GE,
+            ps.EQ,
+            ps.NEQ,
+            ps.InRange,
+            ps.Switch,
+            ps.OR,
+            ps.XOR,
+            ps.AND,
+            ps.Invert,
+            ps.ScalarMaximum,
+            ps.ScalarMinimum,
+            ps.Add,
+            ps.Mul,
+            ps.Sub,
+            ps.TrueDiv,
+            ps.IntDiv,
+            ps.Clip,
+            ps.Second,
+            ps.Identity,
+            ps.Cast,
+            ps.Sign,
+            ps.Neg,
+            ps.Reciprocal,
+            ps.Sqr,
         ]
         scalar_op_amdlibm_speed_up = [
-            aes.Mod,
-            aes.Pow,
-            aes.Ceil,
-            aes.Floor,
-            aes.RoundHalfToEven,
-            aes.RoundHalfAwayFromZero,
-            aes.Log,
-            aes.Log2,
-            aes.Log10,
-            aes.Log1p,
-            aes.Exp,
-            aes.Sqrt,
-            aes.Abs,
-            aes.Cos,
-            aes.Sin,
-            aes.Tan,
-            aes.Tanh,
-            aes.Cosh,
-            aes.Sinh,
-            aes.Sigmoid,
-            aes.Softplus,
+            ps.Mod,
+            ps.Pow,
+            ps.Ceil,
+            ps.Floor,
+            ps.RoundHalfToEven,
+            ps.RoundHalfAwayFromZero,
+            ps.Log,
+            ps.Log2,
+            ps.Log10,
+            ps.Log1p,
+            ps.Exp,
+            ps.Sqrt,
+            ps.Abs,
+            ps.Cos,
+            ps.Sin,
+            ps.Tan,
+            ps.Tanh,
+            ps.Cosh,
+            ps.Sinh,
+            ps.Sigmoid,
+            ps.Softplus,
         ]
 
         def get_scalar_ops(s):
-            if isinstance(s, aes.Composite):
+            if isinstance(s, ps.Composite):
                 l = []
                 for node in s.fgraph.toposort():
                     l += get_scalar_ops(node.op)
@@ -1552,7 +1552,7 @@ def get_scalar_ops(s):
                 return [s]
 
         def list_scalar_op(op):
-            if isinstance(op.scalar_op, aes.Composite):
+            if isinstance(op.scalar_op, ps.Composite):
                 return get_scalar_ops(op.scalar_op)
             else:
                 return [op.scalar_op]
@@ -1579,7 +1579,7 @@ def exp_float32_op(op):
                 return False
             else:
                 l = list_scalar_op(op)
-                return any(s_op.__class__ in [aes.Exp] for s_op in l)
+                return any(s_op.__class__ in [ps.Exp] for s_op in l)
 
         printed_tip = False
         # tip 1

pytensor/ifelse.py

@@ -185,7 +185,7 @@ def make_node(self, condition: "TensorLike", *true_false_branches: Any):
                 input_f.type, HasDataType
             ):
                 # TODO: Be smarter about dtype casting.
-                # up_dtype = aes.upcast(input_t.type.dtype, input_f.type.dtype)
+                # up_dtype = ps.upcast(input_t.type.dtype, input_f.type.dtype)
 
                 if input_t.type.dtype != input_f.type.dtype:
                     raise TypeError(

pytensor/link/jax/dispatch/random.py

@@ -7,7 +7,7 @@
     _coerce_to_uint32_array,
 )
 
-import pytensor.tensor.random.basic as aer
+import pytensor.tensor.random.basic as ptr
 from pytensor.link.jax.dispatch.basic import jax_funcify, jax_typify
 from pytensor.link.jax.dispatch.shape import JAXShapeTuple
 from pytensor.tensor.shape import Shape, Shape_i
@@ -86,7 +86,7 @@ def jax_typify_Generator(rng, **kwargs):
     return state
 
 
-@jax_funcify.register(aer.RandomVariable)
+@jax_funcify.register(ptr.RandomVariable)
 def jax_funcify_RandomVariable(op, node, **kwargs):
     """JAX implementation of random variables."""
     rv = node.outputs[1]
@@ -121,10 +121,10 @@ def jax_sample_fn(op):
     )
 
 
-@jax_sample_fn.register(aer.BetaRV)
-@jax_sample_fn.register(aer.DirichletRV)
-@jax_sample_fn.register(aer.PoissonRV)
-@jax_sample_fn.register(aer.MvNormalRV)
+@jax_sample_fn.register(ptr.BetaRV)
+@jax_sample_fn.register(ptr.DirichletRV)
+@jax_sample_fn.register(ptr.PoissonRV)
+@jax_sample_fn.register(ptr.MvNormalRV)
 def jax_sample_fn_generic(op):
     """Generic JAX implementation of random variables."""
     name = op.name
@@ -140,12 +140,12 @@ def sample_fn(rng, size, dtype, *parameters):
     return sample_fn
 
 
-@jax_sample_fn.register(aer.CauchyRV)
-@jax_sample_fn.register(aer.GumbelRV)
-@jax_sample_fn.register(aer.LaplaceRV)
-@jax_sample_fn.register(aer.LogisticRV)
-@jax_sample_fn.register(aer.NormalRV)
-@jax_sample_fn.register(aer.StandardNormalRV)
+@jax_sample_fn.register(ptr.CauchyRV)
+@jax_sample_fn.register(ptr.GumbelRV)
+@jax_sample_fn.register(ptr.LaplaceRV)
+@jax_sample_fn.register(ptr.LogisticRV)
+@jax_sample_fn.register(ptr.NormalRV)
+@jax_sample_fn.register(ptr.StandardNormalRV)
 def jax_sample_fn_loc_scale(op):
     """JAX implementation of random variables in the loc-scale families.
 
@@ -168,8 +168,8 @@ def sample_fn(rng, size, dtype, *parameters):
     return sample_fn
 
 
-@jax_sample_fn.register(aer.BernoulliRV)
-@jax_sample_fn.register(aer.CategoricalRV)
+@jax_sample_fn.register(ptr.BernoulliRV)
+@jax_sample_fn.register(ptr.CategoricalRV)
 def jax_sample_fn_no_dtype(op):
     """Generic JAX implementation of random variables."""
     name = op.name
@@ -185,9 +185,9 @@ def sample_fn(rng, size, dtype, *parameters):
     return sample_fn
 
 
-@jax_sample_fn.register(aer.RandIntRV)
-@jax_sample_fn.register(aer.IntegersRV)
-@jax_sample_fn.register(aer.UniformRV)
+@jax_sample_fn.register(ptr.RandIntRV)
+@jax_sample_fn.register(ptr.IntegersRV)
+@jax_sample_fn.register(ptr.UniformRV)
 def jax_sample_fn_uniform(op):
     """JAX implementation of random variables with uniform density.
 
@@ -197,7 +197,7 @@ def jax_sample_fn_uniform(op):
     """
     name = op.name
     # IntegersRV is equivalent to RandintRV
-    if isinstance(op, aer.IntegersRV):
+    if isinstance(op, ptr.IntegersRV):
         name = "randint"
     jax_op = getattr(jax.random, name)
 
@@ -214,8 +214,8 @@ def sample_fn(rng, size, dtype, *parameters):
     return sample_fn
 
 
-@jax_sample_fn.register(aer.ParetoRV)
-@jax_sample_fn.register(aer.GammaRV)
+@jax_sample_fn.register(ptr.ParetoRV)
+@jax_sample_fn.register(ptr.GammaRV)
 def jax_sample_fn_shape_scale(op):
     """JAX implementation of random variables in the shape-scale family.
 
@@ -236,7 +236,7 @@ def sample_fn(rng, size, dtype, shape, scale):
     return sample_fn
 
 
-@jax_sample_fn.register(aer.ExponentialRV)
+@jax_sample_fn.register(ptr.ExponentialRV)
 def jax_sample_fn_exponential(op):
     """JAX implementation of `ExponentialRV`."""
 
@@ -251,7 +251,7 @@ def sample_fn(rng, size, dtype, *parameters):
     return sample_fn
 
 
-@jax_sample_fn.register(aer.StudentTRV)
+@jax_sample_fn.register(ptr.StudentTRV)
 def jax_sample_fn_t(op):
     """JAX implementation of `StudentTRV`."""
 
@@ -270,7 +270,7 @@ def sample_fn(rng, size, dtype, *parameters):
     return sample_fn
 
 
-@jax_sample_fn.register(aer.ChoiceRV)
+@jax_sample_fn.register(ptr.ChoiceRV)
 def jax_funcify_choice(op):
     """JAX implementation of `ChoiceRV`."""
 
@@ -285,7 +285,7 @@ def sample_fn(rng, size, dtype, *parameters):
     return sample_fn
 
 
-@jax_sample_fn.register(aer.PermutationRV)
+@jax_sample_fn.register(ptr.PermutationRV)
 def jax_sample_fn_permutation(op):
     """JAX implementation of `PermutationRV`."""
 
@@ -300,7 +300,7 @@ def sample_fn(rng, size, dtype, *parameters):
     return sample_fn
 
 
-@jax_sample_fn.register(aer.BinomialRV)
+@jax_sample_fn.register(ptr.BinomialRV)
 def jax_sample_fn_binomial(op):
     if not numpyro_available:
         raise NotImplementedError(
@@ -323,7 +323,7 @@ def sample_fn(rng, size, dtype, n, p):
     return sample_fn
 
 
-@jax_sample_fn.register(aer.MultinomialRV)
+@jax_sample_fn.register(ptr.MultinomialRV)
 def jax_sample_fn_multinomial(op):
     if not numpyro_available:
         raise NotImplementedError(
@@ -346,7 +346,7 @@ def sample_fn(rng, size, dtype, n, p):
     return sample_fn
 
 
-@jax_sample_fn.register(aer.VonMisesRV)
+@jax_sample_fn.register(ptr.VonMisesRV)
 def jax_sample_fn_vonmises(op):
     if not numpyro_available:
         raise NotImplementedError(

pytensor/link/jax/dispatch/tensor_basic.py

@@ -27,8 +27,8 @@
 ARANGE_CONCRETE_VALUE_ERROR = """JAX requires the arguments of `jax.numpy.arange`
 to be constants. The graph that you defined thus cannot be JIT-compiled
 by JAX. An example of a graph that can be compiled to JAX:
->>> import pytensor.tensor basic
->>> at.arange(1, 10, 2)
+>>> import pytensor.tensor as pt
+>>> pt.arange(1, 10, 2)
 """