Revert optimized->optimize and clarify

Author: ricardoV94

pytensor/link/jax/dispatch/einsum.py

@@ -7,9 +7,8 @@
 @jax_funcify.register(Einsum)
 def jax_funcify_Einsum(op, **kwargs):
     subscripts = op.subscripts
-    optimize = op.optimize
 
     def einsum(*operands):
-        return jnp.einsum(subscripts, *operands, optimize=optimize)
+        return jnp.einsum(subscripts, *operands, optimize="optimal")
 
     return einsum

pytensor/tensor/einsum.py

@@ -33,14 +33,27 @@
 class Einsum(OpFromGraph):
     """
     Wrapper Op for Einsum graphs
+
+    Notes
+    -----
+    The `optimized` prop indicates whether the inner graph was optimized, which can only be done when all shapes are
+    statically known. This is now determined at graph creation time only. We could introduce a rewrite that tries to
+    optimize the graph if static shapes become known later (e.g., after use of `clone_replace` or shape inference during
+    rewrites).
+
+    Also, once the graph is optimized, it could be inlined for potential further optimization that consider the rest of
+    the graph.
+
+    This prop is different from the `optimize` kwarg in numpy that determines what kind (if any) of optimization is
+    desired. We haven't decided whether we want to provide this functionality.
     """
 
-    __props__ = ("subscripts", "path", "optimize")
+    __props__ = ("subscripts", "path", "optimized")
 
-    def __init__(self, *args, subscripts: str, path: str, optimize: bool, **kwargs):
+    def __init__(self, *args, subscripts: str, path: str, optimized: bool, **kwargs):
         self.subscripts = subscripts
         self.path = path
-        self.optimize = optimize
+        self.optimized = optimized
         super().__init__(*args, **kwargs, strict=True)
 
 
@@ -368,7 +381,7 @@ def einsum(subscripts: str, *operands: "TensorLike") -> TensorVariable:
         contraction_list = _contraction_list_from_path(subscripts, operands, path)
 
         # If there are only 1 or 2 operands, there is no optimization to be done?
-        optimize = len(operands) <= 2
+        optimized = len(operands) <= 2
     else:
         # Case 2: All operands have known shapes. In this case, we can use opt_einsum to compute the optimal
         # contraction order.
@@ -382,7 +395,7 @@ def einsum(subscripts: str, *operands: "TensorLike") -> TensorVariable:
             optimize="optimal",
         )
         path = [contraction[0] for contraction in contraction_list]
-        optimize = True
+        optimized = True
 
     def sum_uniques(
         operand: TensorVariable, names: str, uniques: list[str]
@@ -550,6 +563,6 @@ def sum_repeats(
         inputs=list(operands),
         outputs=[einsum_result],
         path=tuple(path),
-        optimize=optimize,
+        optimized=optimized,
     )(*operands)
     return cast(TensorVariable, out)

tests/tensor/test_einsum.py

@@ -130,7 +130,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.optimize == static_shape_known or len(operands) <= 2
+    assert out.owner.op.optimized == static_shape_known or len(operands) <= 2
 
     rng = np.random.default_rng(37)
     test_values = [rng.normal(size=shape).astype(floatX) for shape in shapes]