code: def covariance(
      x, y=None, sample_axis=0, event_axis=-1, keepdims=False, name=None
  ):
      """Sample covariance between observations indexed by `event_axis`.
  
      Given `N` samples of scalar random variables `X` and `Y`, covariance may be
      estimated as
  
      ```none
      Cov[X, Y] := N^{-1} sum_{n=1}^N (X_n - Xbar) Conj{(Y_n - Ybar)}
      Xbar := N^{-1} sum_{n=1}^N X_n
      Ybar := N^{-1} sum_{n=1}^N Y_n
      ```
  
      For vector-variate random variables `X = (X1, ..., Xd)`, `Y = (Y1, ..., Yd)`,
      one is often interested in the covariance matrix, `C_{ij} := Cov[Xi, Yj]`.
  
      ```python
      x = tf.random.normal(shape=(100, 2, 3))
      y = tf.random.normal(shape=(100, 2, 3))
  
      # cov[i, j] is the sample covariance between x[:, i, j] and y[:, i, j].
      cov = tfp.stats.covariance(x, y, sample_axis=0, event_axis=None)
  
      # cov_matrix[i, m, n] is the sample covariance of x[:, i, m] and y[:, i, n]
      cov_matrix = tfp.stats.covariance(x, y, sample_axis=0, event_axis=-1)
      ```
  
      Notice we divide by `N`, which does not create `NaN` when `N = 1`, but is
      slightly biased.
  
      Args:
        x:  A numeric `Tensor` holding samples.
        y:  Optional `Tensor` with same `dtype` and `shape` as `x`.
          Default value: `None` (`y` is effectively set to `x`).
        sample_axis: Scalar or vector `Tensor` designating axis holding samples, or
          `None` (meaning all axis hold samples).
          Default value: `0` (leftmost dimension).
        event_axis:  Scalar or vector `Tensor`, or `None` (scalar events).
          Axis indexing random events, whose covariance we are interested in.
          If a vector, entries must form a contiguous block of dims. `sample_axis`
          and `event_axis` should not intersect.
          Default value: `-1` (rightmost axis holds events).
        keepdims:  Boolean.  Whether to keep the sample axis as singletons.
        name: Python `str` name prefixed to Ops created by this function.
              Default value: `None` (i.e., `'covariance'`).
  
      Returns:
        cov: A `Tensor` of same `dtype` as the `x`, and rank equal to
          `rank(x) - len(sample_axis) + 2 * len(event_axis)`.
  
      Raises:
        AssertionError:  If `x` and `y` are found to have different shape.
        ValueError:  If `sample_axis` and `event_axis` are found to overlap.
        ValueError:  If `event_axis` is found to not be contiguous.
      """
  
      with tf.name_scope(name or "covariance"):
          x = tf.convert_to_tensor(x, name="x")
          # Covariance *only* uses the centered versions of x (and y).
          x = x - tf.reduce_mean(x, axis=sample_axis, keepdims=True)
  
          if y is None:
              y = x
          else:
              y = tf.convert_to_tensor(y, name="y", dtype=x.dtype)
              # If x and y have different shape, sample_axis and event_axis will likely
              # be wrong for one of them!
              tensorshape_util.assert_is_compatible_with(
                  x.shape, y.shape
              )
              y = y - tf.reduce_mean(y, axis=sample_axis, keepdims=True)
  
          if event_axis is None:
              return tf.reduce_mean(
                  x * tf.math.conj(y),
                  axis=sample_axis,
                  keepdims=keepdims,
              )
  
          if sample_axis is None:
              raise ValueError(
                  "sample_axis was None, which means all axis hold events, and this "
                  "overlaps with event_axis ({})".format(event_axis)
              )
  
          event_axis = _make_positive_axis(event_axis, ps.rank(x))
          sample_axis = _make_positive_axis(sample_axis, ps.rank(x))
  
          # If we get lucky and axis is statically defined, we can do some checks.
          if _is_list_like(event_axis) and _is_list_like(sample_axis):
              event_axis = tuple(map(int, event_axis))
              sample_axis = tuple(map(int, sample_axis))
              if set(event_axis).intersection(sample_axis):
                  raise ValueError(
                      "sample_axis ({}) and event_axis ({}) overlapped".format(
                          sample_axis, event_axis
                      )
                  )
              if (np.diff(np.array(sorted(event_axis))) > 1).any():
                  raise ValueError(
                      "event_axis must be contiguous. Found: {}".format(
                          event_axis
                      )
                  )
              batch_axis = list(
                  sorted(
                      set(
                          range(tensorshape_util.rank(x.shape))
                      ).difference(sample_axis + event_axis)
                  )
              )
          else:
              batch_axis = ps.setdiff1d(
                  ps.range(0, ps.rank(x)),
                  ps.concat((sample_axis, event_axis), 0),
              )
  
          event_axis = ps.cast(event_axis, dtype=tf.int32)
          sample_axis = ps.cast(sample_axis, dtype=tf.int32)
          batch_axis = ps.cast(batch_axis, dtype=tf.int32)
  
          # Permute x/y until shape = B + E + S
          perm_for_xy = ps.concat(
              (batch_axis, event_axis, sample_axis), 0
          )
          x_permed = tf.transpose(a=x, perm=perm_for_xy)
          y_permed = tf.transpose(a=y, perm=perm_for_xy)
  
          batch_ndims = ps.size(batch_axis)
          batch_shape = ps.shape(x_permed)[:batch_ndims]
          event_ndims = ps.size(event_axis)
          event_shape = ps.shape(x_permed)[
              batch_ndims : batch_ndims + event_ndims
          ]
          sample_shape = ps.shape(x_permed)[batch_ndims + event_ndims :]
          sample_ndims = ps.size(sample_shape)
          n_samples = ps.reduce_prod(sample_shape)
          n_events = ps.reduce_prod(event_shape)
  
          # Flatten sample_axis into one long dim.
          x_permed_flat = tf.reshape(
              x_permed,
              ps.concat((batch_shape, event_shape, [n_samples]), 0),
          )
          y_permed_flat = tf.reshape(
              y_permed,
              ps.concat((batch_shape, event_shape, [n_samples]), 0),
          )
          # Do the same for event_axis.
          x_permed_flat = tf.reshape(
              x_permed,
              ps.concat((batch_shape, [n_events], [n_samples]), 0),
          )
          y_permed_flat = tf.reshape(
              y_permed,
              ps.concat((batch_shape, [n_events], [n_samples]), 0),
          )
  
          # After matmul, cov.shape = batch_shape + [n_events, n_events]
          cov = tf.matmul(
              x_permed_flat, y_permed_flat, adjoint_b=True
          ) / ps.cast(n_samples, x.dtype)
  
          # Insert some singletons to make
          # cov.shape = batch_shape + event_shape**2 + [1,...,1]
          # This is just like x_permed.shape, except the sample_axis is all 1's, and
          # the [n_events] became event_shape**2.
          cov = tf.reshape(
              cov,
              ps.concat(
                  (
                      batch_shape,
                      # event_shape**2 used here because it is the same length as
                      # event_shape, and has the same number of elements as one
                      # batch of covariance.
                      event_shape**2,
                      ps.ones([sample_ndims], tf.int32),
                  ),
                  0,
              ),
          )
          # Permuting by the argsort inverts the permutation, making
          # cov.shape have ones in the position where there were samples, and
          # [n_events * n_events] in the event position.
          cov = tf.transpose(
              a=cov, perm=ps.invert_permutation(perm_for_xy)
          )
  
          # Now expand event_shape**2 into event_shape + event_shape.
          # We here use (for the first time) the fact that we require event_axis to be
          # contiguous.
          e_start = event_axis[0]
          e_len = 1 + event_axis[-1] - event_axis[0]
          cov = tf.reshape(
              cov,
              ps.concat(
                  (
                      ps.shape(cov)[:e_start],
                      event_shape,
                      event_shape,
                      ps.shape(cov)[e_start + e_len :],
                  ),
                  0,
              ),
          )
  
          # tf.squeeze requires python ints for axis, not Tensor.  This is enough to
          # require our axis args to be constants.
          if not keepdims:
              squeeze_axis = ps.where(
                  sample_axis < e_start,
                  sample_axis,
                  sample_axis + e_len,
              )
              cov = _squeeze(cov, axis=squeeze_axis)
  
          return cov
  
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math.generic import reduce_logmeanexp
label_desc: Use image node_id 3 for calling a global function with example usage: covariance(x, y, sample_axis, event_axis, keepdims, name) and returns: cov, tf