code: def test_psum_scatter(self):
      f = jax.jit(
          jax.pmap(
              lambda x: lax.psum_scatter(x, "i"),
              axis_name="i",
              devices=jax.devices()[:1],
          )
      )
  
      shape = (1, 1, 8)
      x = np.arange(math.prod(shape), dtype=np.float32).reshape(shape)
      self.export_and_compare_to_native(f, x)
  
file_path: jax/tests/export_harnesses_multi_platform_test.py
imports: from __future__ import annotations
  import math
  import re
  from typing import Callable
  from absl import logging
  from absl.testing import absltest
  import numpy
  import jax
  from jax import lax
  from jax._src import test_util
  from jax.experimental import export
  from jax._src.internal_test_util import test_harnesses
label_desc: Use image node_id 4 for calling the PrimitiveTest obj's underlying member method code with example usage: obj.test_psum_scatter() without return types