code: def test_scatter_workaround():
      def make_one_hot_expected(x, i):
          return torch.zeros_like(x).scatter(-1, i.unsqueeze(-1), 1.0)
  
      def make_one_hot_actual(x, i):
          eye = torch.eye(x.shape[-1], dtype=x.dtype, device=x.device)
          return eye[i].clone()
  
      x = torch.randn(5, 4, 3)
      i = torch.randint(0, 3, torch.Size((5, 4)))
      torch.jit.trace(make_one_hot_actual, (x, i))
      expected = make_one_hot_expected(x, i)
      actual = make_one_hot_actual(x, i)
      assert_equal(actual, expected)
  
imports: import logging
  import warnings
  import pytest
  import torch
  from torch.autograd import grad
  from torch.distributions import constraints, kl_divergence
  import pyro
  import pyro.distributions
  import pyro.ops.jit
  import pyro.poutine
  from pyro.distributions.util import scale_and_mask
  from pyro.infer import SVI, JitTrace_ELBO, JitTraceEnum_ELBO, JitTraceGraph_ELBO, JitTraceMeanField_ELBO, Trace_ELBO, TraceEnum_ELBO, TraceGraph_ELBO, TraceMeanField_ELBO, TraceTMC_ELBO, infer_discrete
  from pyro.optim import Adam
  from pyro.poutine.indep_messenger import CondIndepStackFrame
  from pyro.util import ignore_jit_warnings
  from tests.common import assert_close, assert_equal
label_desc: Use image node_id 10 for calling a global function with example usage: test_scatter_workaround() and returns: torch, eye