keep old name 'packed_state' of OptimizerState

BuddenD · Jun 3, 2020 · 0e229e4 · 0e229e4
1 parent b58eec5
commit 0e229e4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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):