def tree_init(self, weight_tree):
    """Assembles node-local initializations into full-tree initialization.

    Args:
      weight_tree: Weights for an entire model, in a tree that matches the
          model's layer structure.

    Returns:
      Tuple `(slots, opt_params)`, where `slots` are the initialized optimizer
      slot values and `opt_params` are optimizer hyperparameters (e.g.,
      learning rate, momentum).
    """
    self._slots = tuple(self.init(weight)
                        for weight in fastmath.tree_flatten(weight_tree))
    return (self._slots, self._init_opt_params)
Assembles node-local initializations into full-tree initialization.

    Args:
      weight_tree: Weights for an entire model, in a tree that matches the
          model's layer structure.

    Returns:
      Tuple `(slots, opt_params)`, where `slots` are the initialized optimizer
      slot values and `opt_params` are optimizer hyperparameters (e.g.,
      learning rate, momentum).
    
trax/optimizers/base.py