make compatible with latest version of optax

rowanz · Jan 25, 2022 · 008c3c1 · 008c3c1
1 parent 8dc049a
commit 008c3c1
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/finetune/optimization.py b/finetune/optimization.py
@@ -12,7 +12,7 @@
 from pretrain.optimization import *
 
 
-class DecayedWeightsDeltaState(OptState):
+class DecayedWeightsDeltaState(NamedTuple):
     """Overall state of the gradient transformation."""
     orig_params: chex.Array  # Momentum
 

diff --git a/pretrain/optimization.py b/pretrain/optimization.py
@@ -1,5 +1,5 @@
 import optax
-from optax import OptState, GradientTransformation
+from optax import GradientTransformation
 from optax._src.base import NO_PARAMS_MSG
 import jax
 import chex
@@ -14,7 +14,7 @@
 from typing import NamedTuple, Any
 
 
-class ScaleByAdamState(OptState):
+class ScaleByAdamState(NamedTuple):
     """State for the Adam algorithm."""
     count: chex.Array
     mu: optax.Updates