Internal change

PiperOrigin-RevId: 428831391

official/modeling/optimization/configs/optimization_config.py

@@ -48,6 +48,8 @@ class OptimizerConfig(oneof.OneOfConfig):
   sgd_experimental: opt_cfg.SGDExperimentalConfig = (
       opt_cfg.SGDExperimentalConfig())
   adam: opt_cfg.AdamConfig = opt_cfg.AdamConfig()
+  adam_experimental: opt_cfg.AdamExperimentalConfig = (
+      opt_cfg.AdamExperimentalConfig())
   adamw: opt_cfg.AdamWeightDecayConfig = opt_cfg.AdamWeightDecayConfig()
   lamb: opt_cfg.LAMBConfig = opt_cfg.LAMBConfig()
   rmsprop: opt_cfg.RMSPropConfig = opt_cfg.RMSPropConfig()

official/modeling/optimization/configs/optimizer_config.py

@@ -67,6 +67,7 @@ class SGDExperimentalConfig(BaseOptimizerConfig):
     name: name of the optimizer.
     nesterov: nesterov for SGD optimizer.
     momentum: momentum for SGD optimizer.
+    jit_compile: if True, jit compile will be used.
   """
   name: str = "SGD"
   nesterov: bool = False
@@ -135,6 +136,30 @@ class AdamConfig(BaseOptimizerConfig):
   amsgrad: bool = False
 
 
+@dataclasses.dataclass
+class AdamExperimentalConfig(BaseOptimizerConfig):
+  """Configuration for experimental Adam optimizer.
+
+  The attributes for this class matches the arguments of
+  `tf.keras.optimizer.experimental.Adam`.
+
+  Attributes:
+    name: name of the optimizer.
+    beta_1: decay rate for 1st order moments.
+    beta_2: decay rate for 2st order moments.
+    epsilon: epsilon value used for numerical stability in Adam optimizer.
+    amsgrad: boolean. Whether to apply AMSGrad variant of this algorithm from
+      the paper "On the Convergence of Adam and beyond".
+    jit_compile: if True, jit compile will be used.
+  """
+  name: str = "Adam"
+  beta_1: float = 0.9
+  beta_2: float = 0.999
+  epsilon: float = 1e-07
+  amsgrad: bool = False
+  jit_compile: bool = False
+
+
 @dataclasses.dataclass
 class AdamWeightDecayConfig(BaseOptimizerConfig):
   """Configuration for Adam optimizer with weight decay.

official/modeling/optimization/optimizer_factory.py

@@ -30,6 +30,7 @@ OPTIMIZERS_CLS = {
     'sgd': tf.keras.optimizers.SGD,
     'sgd_experimental': tf.keras.optimizers.experimental.SGD,
     'adam': tf.keras.optimizers.Adam,
+    'adam_experimental': tf.keras.optimizers.experimental.Adam,
     'adamw': nlp_optimization.AdamWeightDecay,
     'lamb': tfa_optimizers.LAMB,
     'rmsprop': tf.keras.optimizers.RMSprop,