code: def get_moments(x):
      m1 = x.mean(0)
      x = x - m1
      xx = x[..., None] * x[..., None, :]
      m2 = xx.mean(0)
      return torch.cat([m1.reshape(-1), m2.reshape(-1)])
  
file_path: pyro/tests/infer/reparam/test_softmax.py
imports: import pytest
  import torch
  from torch.autograd import grad
  import pyro
  import pyro.distributions
  from pyro import poutine
  from pyro.infer.reparam import GumbelSoftmaxReparam
  from tests.common import assert_close
  from .util import check_init_reparam
label_desc: Use image node_id 1 for calling a global function with example usage: get_moments(x) and returns: torch