jnp.linalg.multi_dot: use optimize='auto'

jax-ml · May 9, 2024 · 5edfaa6 · 5edfaa6
1 parent 1e88e2f
commit 5edfaa6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jax/_src/numpy/linalg.py b/jax/_src/numpy/linalg.py
@@ -2017,4 +2017,4 @@ def multi_dot(arrays: Sequence[ArrayLike], *, precision: PrecisionLike = None) -
   if arrs[-1].ndim == 1:
     einsum_axes[-1] = einsum_axes[-1][:1]
   return jnp.einsum(*itertools.chain(*zip(arrs, einsum_axes)),  # type: ignore[arg-type, call-overload]
-                    optimize='optimal', precision=precision)
+                    optimize='auto', precision=precision)