def print_n_weights(self):
    """Prints the total count of trainable weights."""
    opt_state = self._opt_state
    sizes = _sizes(opt_state.weights)
    if self.n_devices > 1:
      unreplicate = lambda x: x[0]
      single_weights = fastmath.nested_map(unreplicate, opt_state.weights)
      sizes = _sizes(single_weights)
    total_size = _nested_reduce(sum, sizes)
    self.log_step('Total number of trainable weights: %d' % total_size)
Prints the total count of trainable weights.
trax/supervised/trainer_lib.py