class_docstr: Loss balancer.
  
  The loss balancer combines losses together to compute gradients for the backward.
  Given `y = f(...)`, and a number of losses `l1(y, ...)`, `l2(y, ...)`, with `...`
  not having any dependence on `f`, the balancer can efficiently normalize the partial gradients
  `d l1 / d y`, `d l2 / dy` before summing them in order to achieve a desired ratio between
  the losses. For instance if `weights = {'l1': 2, 'l2': 1}`, 66% of the gradient
  going into `f(...)` will come from `l1` on average, and 33% from `l2`. This allows for an easy
  interpration of the weights even if the intrisic scale of `l1`, `l2` ... is unknown.
  
  Noting `g1 = d l1 / dy`, etc., the balanced gradient `G` will be
  (with `avg` an exponential moving average over the updates),
  
      G = sum_i total_norm * g_i / avg(||g_i||) * w_i / sum(w_i)
  
  If `balance_grads` is False, this is deactivated, and instead the gradient will just be the
  standard sum of the partial gradients with the given weights.
  
  A call to the backward method of the balancer will compute the the partial gradients,
  combining all the losses and potentially rescaling the gradients,
  which can help stabilize the training and reason about multiple losses with varying scales.
  The obtained gradient with respect to `y` is then back-propagated to `f(...)`.
  
  Expected usage:
  
      weights = {'loss_a': 1, 'loss_b': 4}
      balancer = Balancer(weights, ...)
      losses: dict = {}
      losses['loss_a'] = compute_loss_a(x, y)
      losses['loss_b'] = compute_loss_b(x, y)
      if model.training():
          effective_loss = balancer.backward(losses, x)
  
  Args:
      weights (dict[str, float]): Weight coefficient for each loss. The balancer expect the losses keys
          from the backward method to match the weights keys to assign weight to each of the provided loss.
      balance_grads (bool): Whether to rescale gradients so that weights reflect the fraction of the
          overall gradient, rather than a constant multiplier.
      total_norm (float): Reference norm when rescaling gradients, ignored otherwise.
      emay_decay (float): EMA decay for averaging the norms.
      per_batch_item (bool): Whether to compute the averaged norm per batch item or not. This only holds
          when rescaling the gradients.
      epsilon (float): Epsilon value for numerical stability.
      monitor (bool): If True, stores in `self.metrics` the relative ratio between the norm of the gradients
          coming from each loss, when calling `backward()`.
code: def __init__(
      self,
      weights: tp.Dict[str, float],
      balance_grads: bool = True,
      total_norm: float = 1.0,
      ema_decay: float = 0.999,
      per_batch_item: bool = True,
      epsilon: float = 1e-12,
      monitor: bool = False,
  ):
      self.weights = weights
      self.per_batch_item = per_batch_item
      self.total_norm = total_norm or 1.0
      self.averager = flashy.averager(ema_decay or 1.0)
      self.epsilon = epsilon
      self.monitor = monitor
      self.balance_grads = balance_grads
      self._metrics: tp.Dict[str, tp.Any] = {}
  
file_path: audiocraft/audiocraft/losses/balancer.py
imports: import typing
  import flashy
  import torch
  from torch import autograd
label_desc: Use image node_id 1 to create a new Balancer object  with example: obj = Balancer(weights, balance_grads, total_norm, ema_decay, per_batch_item, epsilon, monitor)