From 0e229e4ea8c7361e1a72c2b15b9f18d9829c0a03 Mon Sep 17 00:00:00 2001
From: Matthew Johnson <mattjj@google.com>
Date: Wed, 3 Jun 2020 07:32:44 -0700
Subject: [PATCH] keep old name 'packed_state' of OptimizerState

---
 jax/experimental/optimizers.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/jax/experimental/optimizers.py b/jax/experimental/optimizers.py
index d1449329ed18..2b92ae24bb49 100644
--- a/jax/experimental/optimizers.py
+++ b/jax/experimental/optimizers.py
@@ -89,10 +89,10 @@
 # lists (with no further nesting).
 
 OptimizerState = namedtuple("OptimizerState",
-                            ["states_flat", "tree_def", "subtree_defs"])
+                            ["packed_state", "tree_def", "subtree_defs"])
 register_pytree_node(
     OptimizerState,
-    lambda xs: ((xs.states_flat,), (xs.tree_def, xs.subtree_defs)),
+    lambda xs: ((xs.packed_state,), (xs.tree_def, xs.subtree_defs)),
     lambda data, xs: OptimizerState(xs[0], data[0], data[1]))
 
 def optimizer(opt_maker):