class_docstr: Optimizer that implements the Layer-wise Adaptive Moments (LAMB).
  
  See paper [Large Batch Optimization for Deep Learning: Training BERT
  in 76 minutes](https://arxiv.org/abs/1904.00962).
class_docstr_tok: ["Optimizer","that","implements","the","Layer-wise","Adaptive","Moments","(","LAMB",")",".","See","paper","[","Large","Batch","Optimization","for","Deep","Learning",":","Training","BERT","in","76","minutes","]","(","https",":","\/\/arxiv.org\/abs\/1904.00962",")","."]
code: def _prepare_local(self, var_device, var_dtype, apply_state):
      super()._prepare_local(var_device, var_dtype, apply_state)
  
      local_step = tf.cast(self.iterations + 1, var_dtype)
      beta_1_t = tf.identity(self._get_hyper("beta_1", var_dtype))
      beta_2_t = tf.identity(self._get_hyper("beta_2", var_dtype))
      weight_decay = tf.identity(
          self._get_hyper("weight_decay", var_dtype)
      )
      beta_1_power = tf.pow(beta_1_t, local_step)
      beta_2_power = tf.pow(beta_2_t, local_step)
      apply_state[(var_device, var_dtype)].update(
          dict(
              weight_decay=weight_decay,
              epsilon=tf.convert_to_tensor(self.epsilon, var_dtype),
              beta_1_t=beta_1_t,
              beta_1_power=beta_1_power,
              one_minus_beta_1_t=1 - beta_1_t,
              beta_2_t=beta_2_t,
              beta_2_power=beta_2_power,
              one_minus_beta_2_t=1 - beta_2_t,
          )
      )
  
file_path: addons/tensorflow_addons/optimizers/lamb.py
imports: import warnings
  from typing import Optional, Union, Callable, List
  from typeguard import typechecked
  import tensorflow
  from tensorflow_addons.optimizers import KerasLegacyOptimizer
  from tensorflow_addons.utils.types import FloatTensorLike
  from tensorflow_addons.optimizers.utils import is_variable_matched_by_regexes
label_desc: Use image node_id 3 for calling the LAMB obj's underlying member method code with example usage: obj._prepare_local(var_device, var_dtype, apply_state) without return types