code: def gelman_rubin(input, chain_dim=0, sample_dim=1):
      """
      Computes R-hat over chains of samples. It is required that
      ``input.size(sample_dim) >= 2`` and ``input.size(chain_dim) >= 2``.
  
      :param torch.Tensor input: the input tensor.
      :param int chain_dim: the chain dimension.
      :param int sample_dim: the sample dimension.
      :returns torch.Tensor: R-hat of ``input``.
      """
      assert input.dim() >= 2
      assert input.size(sample_dim) >= 2
      assert input.size(chain_dim) >= 2
      # change input.shape to 1 x 1 x input.shape
      # then transpose sample_dim with 0, chain_dim with 1
      sample_dim = (
          input.dim() + sample_dim if sample_dim < 0 else sample_dim
      )
      chain_dim = (
          input.dim() + chain_dim if chain_dim < 0 else chain_dim
      )
      assert chain_dim != sample_dim
      input = input.reshape((1, 1) + input.shape)
      input = input.transpose(0, sample_dim + 2).transpose(
          1, chain_dim + 2
      )
  
      var_within, var_estimator = _compute_chain_variance_stats(input)
      rhat = (var_estimator / var_within).sqrt()
      return rhat.squeeze(max(sample_dim, chain_dim)).squeeze(
          min(sample_dim, chain_dim)
      )
  
imports: import math
  import numbers
  import torch
  from torch.fft import irfft, rfft
  from .tensor_utils import next_fast_len
label_desc: Use image node_id 2 for calling a global function with example usage: gelman_rubin(input, chain_dim, sample_dim) and returns: rhat