docs: def submit(self, value):
      """Submit a single or batch tensor to refine the streaming mean."""
      # Add a batch dimension if necessary.
      if value.shape.ndims == self._sum.shape.ndims:
        value = value[None, ...]
      return tf.group(
          self._sum.assign_add(tf.reduce_sum(value, 0)),
          self._count.assign_add(tf.shape(value)[0]))