try optimize=True with einsum

closes #2583 can revert if this ends up problematic for some reason!
jax-ml · Apr 21, 2020 · 783e24f · 783e24f
1 parent 18f9674
commit 783e24f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
@@ -2558,7 +2558,7 @@ def tensordot(a, b, axes=2, precision=None):
 
 @_wraps(onp.einsum, lax_description=_PRECISION_DOC)
 def einsum(*operands, **kwargs):
-  optimize = kwargs.pop('optimize', 'auto')
+  optimize = kwargs.pop('optimize', True)
   optimize = 'greedy' if optimize is True else optimize
   precision = kwargs.pop('precision', None)
   if kwargs: