class_docstr: Layerwise adaptive rate scaling (LARS) optimizer.
  
  See https://arxiv.org/abs/1708.03888
class_docstr_tok: ["Layerwise","adaptive","rate","scaling","(","LARS",")","optimizer",".","See","https",":","\/\/arxiv.org\/abs\/1708.03888"]
docstr: Updates variables based on LARS algorithm.
  
  Args:
      lr: learning rate. The LARS paper suggests using lr = lr_0 * (1 -t/T)**2,
      where t is the current epoch number and T the maximum number of epochs.
      grads: the gradients to apply.
docstr_tok: ["Updates","variables","based","on","LARS","algorithm",".","Args",":","lr",":","learning","rate",".","The","LARS","paper","suggests","using","lr","=","lr_0","*","(","1","-t\/T",")","*","*","2",",","where","t","is","the","current","epoch","number","and","T","the","maximum","number","of","epochs",".","grads",":","the","gradients","to","apply","."]
code: def __call__(self, lr: float, grads: List[JaxArray]):
      """Updates variables based on LARS algorithm.
  
      Args:
          lr: learning rate. The LARS paper suggests using lr = lr_0 * (1 -t/T)**2,
          where t is the current epoch number and T the maximum number of epochs.
          grads: the gradients to apply.
      """
      assert len(grads) == len(
          self.train_vars
      ), "Expecting as many gradients as trainable variables"
  
      for g, p, m in zip(grads, self.train_vars, self.m):
          p_norm = jn.linalg.norm(p.value)
          g_norm = jn.linalg.norm(g)
          trust_ratio = (
              self.tc
              * p_norm
              / (g_norm + self.weight_decay * p_norm + self.eps)
          )
          local_lr = lr * jn.maximum(
              jn.logical_or(p_norm == 0, g_norm == 0), trust_ratio
          )
          m.value = self.momentum * m.value + local_lr * (
              g + self.weight_decay * p.value
          )
          p.value -= m.value
  
file_path: objax/objax/optimizer/lars.py
imports: from typing import List
  import jax.numpy
  from objax.module import Module, ModuleList
  from objax.typing import JaxArray
  from objax.util import class_name
  from objax.variable import TrainRef, StateVar, TrainVar, VarCollection
label_desc: Use image node_id 2 for calling the LARS obj's underlying member method code with example usage: obj.__call__(lr, grads) without return types