code: def _potential_scale_reduction_single_state(
      state, independent_chain_ndims, split_chains, validate_args
  ):
      """potential_scale_reduction for one single state `Tensor`."""
      # casting integers to floats for floating-point division
      # check to see if the `state` is a numpy object for the numpy test suite
      if dtype_util.as_numpy_dtype(state.dtype) is np.int64:
          state = tf.cast(state, tf.float64)
      elif dtype_util.is_integer(state.dtype):
          state = tf.cast(state, tf.float32)
      with tf.name_scope("potential_scale_reduction_single_state"):
          # We assume exactly one leading dimension indexes e.g. correlated samples
          # from each Markov chain.
          state = tf.convert_to_tensor(state, name="state")
  
          n_samples_ = tf.compat.dimension_value(state.shape[0])
          if n_samples_ is not None:  # If available statically.
              if split_chains and n_samples_ < 4:
                  raise ValueError(
                      "Must provide at least 4 samples when splitting chains. "
                      "Found {}".format(n_samples_)
                  )
              if not split_chains and n_samples_ < 2:
                  raise ValueError(
                      "Must provide at least 2 samples.  Found {}".format(
                          n_samples_
                      )
                  )
          elif validate_args:
              if split_chains:
                  assertions = [
                      assert_util.assert_greater(
                          ps.shape(state)[0],
                          4,
                          message="Must provide at least 4 samples when splitting chains.",
                      )
                  ]
                  with tf.control_dependencies(assertions):
                      state = tf.identity(state)
              else:
                  assertions = [
                      assert_util.assert_greater(
                          ps.shape(state)[0],
                          2,
                          message="Must provide at least 2 samples.",
                      )
                  ]
                  with tf.control_dependencies(assertions):
                      state = tf.identity(state)
  
          # Define so it's not a magic number.
          # Warning!  `if split_chains` logic assumes this is 1!
          sample_ndims = 1
  
          if split_chains:
              # Split the sample dimension in half, doubling the number of
              # independent chains.
  
              # For odd number of samples, keep all but the last sample.
              state_shape = ps.shape(state)
              n_samples = state_shape[0]
              state = state[: n_samples - n_samples % 2]
  
              # Suppose state = [0, 1, 2, 3, 4, 5]
              # Step 1: reshape into [[0, 1, 2], [3, 4, 5]]
              # E.g. reshape states of shape [a, b] into [2, a//2, b].
              state = tf.reshape(
                  state,
                  ps.concat(
                      [[2, n_samples // 2], state_shape[1:]], axis=0
                  ),
              )
              # Step 2: Put the size `2` dimension in the right place to be treated as a
              # chain, changing [[0, 1, 2], [3, 4, 5]] into [[0, 3], [1, 4], [2, 5]],
              # reshaping [2, a//2, b] into [a//2, 2, b].
              state = tf.transpose(
                  a=state,
                  perm=ps.concat(
                      [[1, 0], ps.range(2, ps.rank(state))], axis=0
                  ),
              )
  
              # We're treating the new dim as indexing 2 chains, so increment.
              independent_chain_ndims += 1
  
          sample_axis = ps.range(0, sample_ndims)
          chain_axis = ps.range(
              sample_ndims, sample_ndims + independent_chain_ndims
          )
          sample_and_chain_axis = ps.range(
              0, sample_ndims + independent_chain_ndims
          )
  
          n = _axis_size(state, sample_axis)
          m = _axis_size(state, chain_axis)
  
          # In the language of Brooks and Gelman (1998),
          # B / n is the between chain variance, the variance of the chain means.
          # W is the within sequence variance, the mean of the chain variances.
          b_div_n = _reduce_variance(
              tf.reduce_mean(state, axis=sample_axis, keepdims=True),
              sample_and_chain_axis,
              biased=False,
          )
          w = tf.reduce_mean(
              _reduce_variance(
                  state, sample_axis, keepdims=True, biased=False
              ),
              axis=sample_and_chain_axis,
          )
  
          # sigma^2_+ is an estimate of the true variance, which would be unbiased if
          # each chain was drawn from the target.  c.f. "law of total variance."
          sigma_2_plus = ((n - 1) / n) * w + b_div_n
          return ((m + 1.0) / m) * sigma_2_plus / w - (n - 1.0) / (
              m * n
          )
  
file_path: probability/tensorflow_probability/python/mcmc/diagnostic.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python import stats
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import nest_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow.python.util import nest
label_desc: Use image node_id 4 for calling a global function with example usage: _potential_scale_reduction_single_state(state, independent_chain_ndims, split_chains, validate_args) and returns: unknown