code: def test_tensor_distribution(event_inputs, batch_inputs, test_grad):
      num_samples = 50000
      sample_inputs = OrderedDict(n=Bint[num_samples])
      be_inputs = OrderedDict(batch_inputs + event_inputs)
      batch_inputs = OrderedDict(batch_inputs)
      event_inputs = OrderedDict(event_inputs)
      sampled_vars = frozenset(event_inputs)
      p_data = random_tensor(be_inputs).data
      rng_key = (
          None
          if get_backend() == "torch"
          else np.array([0, 0], dtype=np.uint32)
      )
      probe = randn(p_data.shape)
  
      def diff_fn(p_data):
          p = Tensor(p_data, be_inputs)
          q = p.sample(sampled_vars, sample_inputs, rng_key=rng_key)
          q -= math.log(num_samples)
          mq = p.materialize(q).reduce(ops.logaddexp, "n")
          mq = mq.align(tuple(p.inputs))
  
          _, (p_data, mq_data) = align_tensors(p, mq)
          assert p_data.shape == mq_data.shape
          return (ops.exp(mq_data) * probe).sum() - (
              ops.exp(p_data) * probe
          ).sum(), mq
  
      if test_grad:
          if get_backend() == "jax":
              import jax
  
              diff_grad, mq = jax.grad(diff_fn, has_aux=True)(p_data)
          else:
              import torch
  
              p_data.requires_grad_(True)
              diff_grad = torch.autograd.grad(
                  diff_fn(p_data)[0], [p_data]
              )[0]
  
          assert_close(
              diff_grad,
              ops.new_zeros(diff_grad, diff_grad.shape),
              atol=0.1,
              rtol=None,
          )
      else:
          _, mq = diff_fn(p_data)
          assert_close(
              mq, Tensor(p_data, be_inputs), atol=0.1, rtol=None
          )
  
imports: import itertools
  import math
  from collections import OrderedDict
  from importlib import import_module
  import numpy
  import pytest
  import funsor.ops
  from funsor.cnf import Contraction
  from funsor.delta import Delta
  from funsor.distribution import BACKEND_TO_DISTRIBUTIONS_BACKEND
  from funsor.domains import Bint, Real, Reals
  from funsor.integrate import Integrate
  from funsor.montecarlo import MonteCarlo
  from funsor.tensor import Tensor, align_tensors
  from funsor.terms import Variable
  from funsor.testing import assert_close, excludes_backend, id_from_inputs, randn, random_gaussian, random_tensor, xfail_if_not_implemented
  from funsor.util import get_backend
label_desc: Use image node_id 5 for calling a global function with example usage: test_tensor_distribution(event_inputs, batch_inputs, test_grad) and returns: unknown,  mq