code: def _moments(
      x, axes, shift=None, keepdims=False, name=None
  ):  # pylint: disable=unused-argument
      # NOTE: If x.dtype is float16, we may want to compute in float32.
      mean = reduce_mean(x, axis=axes, keepdims=True)
      # NOTE: The gradient backpropagated to the mean from the variance calcuation
      # is zero, so we can safely use `stop_gradient(mean)` for efficiency.
      variance = reduce_mean(
          squared_difference(x, stop_gradient(mean)),
          axis=axes,
          keepdims=keepdims,
      )
      if not keepdims:
          mean = numpy_array.squeeze(mean, axes)
      return (mean, variance)
  
file_path: probability/tensorflow_probability/python/internal/backend/numpy/nn.py
imports: import collections
  import numpy
  from tensorflow_probability.python.internal.backend.numpy import _utils
  from tensorflow_probability.python.internal.backend.numpy import numpy_array
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import l2_normalize
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import log_softmax
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import reduce_logsumexp
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import reduce_mean
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import softmax
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import softplus
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import squared_difference
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import top_k
  from tensorflow_probability.python.internal.backend.numpy.ops import stop_gradient
label_desc: Use image node_id 6 for calling a global function with example usage: _moments(x, axes, shift, keepdims, name) and returns: mean,  variance