code: def marginal_prob(prob, axis):
      """Compute the marginal probability given a joint probability distribution expressed as a tensor.
      Each random variable corresponds to a dimension.
  
      If the distribution arises from a quantum circuit measured in computational basis, each dimension
      corresponds to a wire. For example, for a 2-qubit quantum circuit `prob[0, 1]` is the probability of measuring the
      first qubit in state 0 and the second in state 1.
  
      Args:
          prob (tensor_like): 1D tensor of probabilities. This tensor should of size
              ``(2**N,)`` for some integer value ``N``.
          axis (list[int]): the axis for which to calculate the marginal
              probability distribution
  
      Returns:
          tensor_like: the marginal probabilities, of
          size ``(2**len(axis),)``
  
      **Example**
  
      >>> x = tf.Variable([1, 0, 0, 1.], dtype=tf.float64) / np.sqrt(2)
      >>> marginal_prob(x, axis=[0, 1])
      <tf.Tensor: shape=(4,), dtype=float64, numpy=array([0.70710678, 0.        , 0.        , 0.70710678])>
      >>> marginal_prob(x, axis=[0])
      <tf.Tensor: shape=(2,), dtype=float64, numpy=array([0.70710678, 0.70710678])>
      """
      prob = np.flatten(prob)
      num_wires = int(np.log2(len(prob)))
  
      if num_wires == len(axis):
          return prob
  
      inactive_wires = tuple(set(range(num_wires)) - set(axis))
      prob = np.reshape(prob, [2] * num_wires)
      prob = np.sum(prob, axis=inactive_wires)
      return np.flatten(prob)
  
imports: import itertools
  import functools
  from string import ascii_letters
  from autoray import numpy
  from numpy import float64
  import pennylane
  from .None import single_dispatch
  from .multi_dispatch import diag, dot, scatter_element_add, einsum, get_interface
  from .utils import is_abstract, allclose, cast, convert_like, cast_like
  from .matrix_manipulation import _permute_dense_matrix
label_desc: Use image node_id 2 for calling a global function with example usage: marginal_prob(prob, axis) and returns: np, prob