Rename optimized->optimize to match numpy/jax signature

Author: jessegrabowski

pytensor/tensor/einsum.py

@@ -35,12 +35,12 @@ class Einsum(OpFromGraph):
     Wrapper Op for Einsum graphs
     """
 
-    __props__ = ("subscripts", "path", "optimized")
+    __props__ = ("subscripts", "path", "optimize")
 
-    def __init__(self, *args, subscripts: str, path: str, optimized: bool, **kwargs):
+    def __init__(self, *args, subscripts: str, path: str, optimize: bool, **kwargs):
         self.subscripts = subscripts
         self.path = path
-        self.optimized = optimized
+        self.optimize = optimize
         super().__init__(*args, **kwargs, strict=True)
 
 
@@ -223,7 +223,7 @@ def einsum(subscripts: str, *operands: "TensorLike") -> TensorVariable:
     shapes = [operand.type.shape for operand in operands]
 
     if None in itertools.chain.from_iterable(shapes):
-        # We mark optimized = False, even in cases where there is no ordering optimization to be done
+        # We mark optimize = False, even in cases where there is no ordering optimization to be done
         # because the inner graph may have to accommodate dynamic shapes.
         # If those shapes become known later we will likely want to rebuild the Op (unless we inline it)
         if len(operands) == 1:
@@ -234,7 +234,7 @@ def einsum(subscripts: str, *operands: "TensorLike") -> TensorVariable:
             # We use (1,0) and not (0,1) because that's what opt_einsum tends to prefer, and so the Op signatures will match more often
             path = [(1, 0) for i in range(len(operands) - 1)]
         contraction_list = contraction_list_from_path(subscripts, operands, path)
-        optimized = (
+        optimize = (
             len(operands) <= 2
         )  # If there are only 1 or 2 operands, there is no optimization to be done?
     else:
@@ -247,7 +247,7 @@ def einsum(subscripts: str, *operands: "TensorLike") -> TensorVariable:
             shapes=True,
         )
         path = [contraction[0] for contraction in contraction_list]
-        optimized = True
+        optimize = True
 
     def sum_uniques(
         operand: TensorVariable, names: str, uniques: list[str]
@@ -412,6 +412,6 @@ def sum_repeats(
         inputs=list(operands),
         outputs=[einsum_result],
         path=tuple(path),
-        optimized=optimized,
+        optimize=optimize,
     )(*operands)
     return cast(TensorVariable, out)

tests/tensor/test_einsum.py

@@ -127,7 +127,7 @@ def test_einsum_signatures(static_shape_known, signature):
         for name, static_shape in zip(ascii_lowercase, static_shapes)
     ]
     out = pt.einsum(signature, *operands)
-    assert out.owner.op.optimized == static_shape_known or len(operands) <= 2
+    assert out.owner.op.optimize == static_shape_known or len(operands) <= 2
 
     rng = np.random.default_rng(37)
     test_values = [rng.normal(size=shape) for shape in shapes]