def __init__(self, name, aggregation_fun=jnp.mean):
    """Takes a tensor and returns it.

    Args:
      name: Name of the metric to be reported.
      aggregation_fun: Aggregation function to be used.
    """
    super().__init__(name=f'Summary_{name}', n_in=1, n_out=1)
    name = 'summary_' + name
    self._name = name
    self._aggregation_fun = aggregation_fun
Takes a tensor and returns it.

    Args:
      name: Name of the metric to be reported.
      aggregation_fun: Aggregation function to be used.
    
trax/layers/core.py