code: def auto_correlation(
      x,
      axis=-1,
      max_lags=None,
      center=True,
      normalize=True,
      name="auto_correlation",
  ):
      """Auto correlation along one axis.
  
      Given a `1-D` wide sense stationary (WSS) sequence `X`, the auto correlation
      `RXX` may be defined as  (with `E` expectation and `Conj` complex conjugate)
  
      ```
      RXX[m] := E{ W[m] Conj(W[0]) } = E{ W[0] Conj(W[-m]) },
      W[n]   := (X[n] - MU) / S,
      MU     := E{ X[0] },
      S**2   := E{ (X[0] - MU) Conj(X[0] - MU) }.
      ```
  
      This function takes the viewpoint that `x` is (along one axis) a finite
      sub-sequence of a realization of (WSS) `X`, and then uses `x` to produce an
      estimate of `RXX[m]` as follows:
  
      After extending `x` from length `L` to `inf` by zero padding, the auto
      correlation estimate `rxx[m]` is computed for `m = 0, 1, ..., max_lags` as
  
      ```
      rxx[m] := (L - m)**-1 sum_n w[n + m] Conj(w[n]),
      w[n]   := (x[n] - mu) / s,
      mu     := L**-1 sum_n x[n],
      s**2   := L**-1 sum_n (x[n] - mu) Conj(x[n] - mu)
      ```
  
      The error in this estimate is proportional to `1 / sqrt(len(x) - m)`, so users
      often set `max_lags` small enough so that the entire output is meaningful.
  
      Note that since `mu` is an imperfect estimate of `E{ X[0] }`, and we divide by
      `len(x) - m` rather than `len(x) - m - 1`, our estimate of auto correlation
      contains a slight bias, which goes to zero as `len(x) - m --> infinity`.
  
      Args:
        x:  `float32` or `complex64` `Tensor`.
        axis:  Python `int`. The axis number along which to compute correlation.
          Other dimensions index different batch members.
        max_lags:  Positive `int` tensor.  The maximum value of `m` to consider (in
          equation above).  If `max_lags >= x.shape[axis]`, we effectively re-set
          `max_lags` to `x.shape[axis] - 1`.
        center:  Python `bool`.  If `False`, do not subtract the mean estimate `mu`
          from `x[n]` when forming `w[n]`.
        normalize:  Python `bool`.  If `False`, do not divide by the variance
          estimate `s**2` when forming `w[n]`.
        name:  `String` name to prepend to created ops.
  
      Returns:
        `rxx`: `Tensor` of same `dtype` as `x`.  `rxx.shape[i] = x.shape[i]` for
          `i != axis`, and `rxx.shape[axis] = max_lags + 1`.
  
      Raises:
        TypeError:  If `x` is not a supported type.
      """
      # Implementation details:
      # Extend length N / 2 1-D array x to length N by zero padding onto the end.
      # Then, set
      #   F[x]_k := sum_n x_n exp{-i 2 pi k n / N }.
      # It is not hard to see that
      #   F[x]_k Conj(F[x]_k) = F[R]_k, where
      #   R_m := sum_n x_n Conj(x_{(n - m) mod N}).
      # One can also check that R_m / (N / 2 - m) is an unbiased estimate of RXX[m].
  
      # Since F[x] is the DFT of x, this leads us to a zero-padding and FFT/IFFT
      # based version of estimating RXX.
      # Note that this is a special case of the Wiener-Khinchin Theorem.
      with tf.name_scope(name):
          x = tf.convert_to_tensor(x, name="x")
  
          # Rotate dimensions of x in order to put axis at the rightmost dim.
          # FFT op requires this.
          rank = ps.rank(x)
          if axis < 0:
              axis = rank + axis
          shift = rank - 1 - axis
          # Suppose x.shape[axis] = T, so there are T 'time' steps.
          #   ==> x_rotated.shape = B + [T],
          # where B is x_rotated's batch shape.
          x_rotated = distribution_util.rotate_transpose(x, shift)
  
          if center:
              x_rotated = x_rotated - tf.reduce_mean(
                  x_rotated, axis=-1, keepdims=True
              )
  
          # x_len = N / 2 from above explanation.  The length of x along axis.
          # Get a value for x_len that works in all cases.
          x_len = ps.shape(x_rotated)[-1]
  
          # TODO(langmore) Investigate whether this zero padding helps or hurts.  At
          # the moment is necessary so that all FFT implementations work.
          # Zero pad to the next power of 2 greater than 2 * x_len, which equals
          # 2**(ceil(Log_2(2 * x_len))).  Note: Log_2(X) = Log_e(X) / Log_e(2).
          x_len_float64 = ps.cast(x_len, np.float64)
          target_length = ps.pow(
              np.float64(2.0),
              ps.ceil(ps.log(x_len_float64 * 2) / np.log(2.0)),
          )
          pad_length = ps.cast(target_length - x_len_float64, np.int32)
  
          # We should have:
          # x_rotated_pad.shape = x_rotated.shape[:-1] + [T + pad_length]
          #                     = B + [T + pad_length]
          x_rotated_pad = distribution_util.pad(
              x_rotated, axis=-1, back=True, count=pad_length
          )
  
          dtype = x.dtype
          if not dtype_util.is_complex(dtype):
              if not dtype_util.is_floating(dtype):
                  raise TypeError(
                      "Argument x must have either float or complex dtype"
                      " found: {}".format(dtype)
                  )
              x_rotated_pad = tf.complex(
                  x_rotated_pad,
                  dtype_util.as_numpy_dtype(
                      dtype_util.real_dtype(dtype)
                  )(0.0),
              )
  
          # Autocorrelation is IFFT of power-spectral density (up to some scaling).
          fft_x_rotated_pad = tf.signal.fft(x_rotated_pad)
          spectral_density = fft_x_rotated_pad * tf.math.conj(
              fft_x_rotated_pad
          )
          # shifted_product is R[m] from above detailed explanation.
          # It is the inner product sum_n X[n] * Conj(X[n - m]).
          shifted_product = tf.signal.ifft(spectral_density)
  
          # Cast back to real-valued if x was real to begin with.
          shifted_product = tf.cast(shifted_product, dtype)
  
          # Figure out if we can deduce the final static shape, and set max_lags.
          # Use x_rotated as a reference, because it has the time dimension in the far
          # right, and was created before we performed all sorts of crazy shape
          # manipulations.
          know_static_shape = True
          if not tensorshape_util.is_fully_defined(x_rotated.shape):
              know_static_shape = False
          if max_lags is None:
              max_lags = x_len - 1
          else:
              max_lags = ps.convert_to_shape_tensor(
                  max_lags, name="max_lags"
              )
              max_lags_ = tf.get_static_value(max_lags)
              if max_lags_ is None or not know_static_shape:
                  know_static_shape = False
                  max_lags = tf.minimum(x_len - 1, max_lags)
              else:
                  max_lags = min(x_len - 1, max_lags_)
  
          # Chop off the padding.
          # We allow users to provide a huge max_lags, but cut it off here.
          # shifted_product_chopped.shape = x_rotated.shape[:-1] + [max_lags]
          shifted_product_chopped = shifted_product[..., : max_lags + 1]
  
          # If possible, set shape.
          if know_static_shape:
              chopped_shape = tensorshape_util.as_list(x_rotated.shape)
              chopped_shape[-1] = min(x_len, max_lags + 1)
              tensorshape_util.set_shape(
                  shifted_product_chopped, chopped_shape
              )
  
          # Recall R[m] is a sum of N / 2 - m nonzero terms x[n] Conj(x[n - m]).  The
          # other terms were zeros arising only due to zero padding.
          # `denominator = (N / 2 - m)` (defined below) is the proper term to
          # divide by to make this an unbiased estimate of the expectation
          # E[X[n] Conj(X[n - m])].
          x_len = ps.cast(x_len, dtype_util.real_dtype(dtype))
          max_lags = ps.cast(max_lags, dtype_util.real_dtype(dtype))
          denominator = x_len - ps.range(0.0, max_lags + 1.0)
          denominator = ps.cast(denominator, dtype)
          shifted_product_rotated = (
              shifted_product_chopped / denominator
          )
  
          if normalize:
              shifted_product_rotated /= shifted_product_rotated[
                  ..., :1
              ]
  
          # Transpose dimensions back to those of x.
          return distribution_util.rotate_transpose(
              shifted_product_rotated, -shift
          )
  
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 1 for calling a global function with example usage: auto_correlation(x, axis, max_lags, center, normalize, name) and returns: distribution_util