From 783e24fd4557c35254772f6a255f30b5bba69b67 Mon Sep 17 00:00:00 2001
From: Matthew Johnson <mattjj@google.com>
Date: Tue, 21 Apr 2020 15:01:30 -0700
Subject: [PATCH] try optimize=True with einsum

closes #2583

can revert if this ends up problematic for some reason!
---
 jax/numpy/lax_numpy.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
index b4ced0dfaed0..df61fded92bf 100644
--- 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: