class_docstr: Test simple_gather().
class_docstr_tok: ["Test","simple_gather","(",")","."]
docstr: Test against manually specified target when idx is a nd-array.
docstr_tok: ["Test","against","manually","specified","target","when","idx","is","a","nd-array","."]
code: def test_simple_gather_ndarray(self):
      """Test against manually specified target when idx is a nd-array."""
      x = jnp.array(np.random.normal(size=(2, 3, 5)), dtype=jnp.float32)
      idx = jnp.array([[1, 0, 2], [2, 1, 0]], dtype=jnp.int32)
      y = model_utils.simple_gather(x, idx)
      y_target = jnp.stack(
          [
              jnp.stack([x[0, 1], x[0, 0], x[0, 2]]),
              jnp.stack([x[1, 2], x[1, 1], x[1, 0]]),
          ]
      )
  
      self.assertSequenceAlmostEqual(y.flatten(), y_target.flatten())
  
imports: import itertools
  from absl.testing import absltest
  from absl.testing import parameterized
  from flax.training import common_utils
  import jax
  import jax.numpy
  import numpy
  from scenic.model_lib.base_models import model_utils
label_desc: Use image node_id 1 for calling the SimpleGatherTest obj's underlying member method code with example usage: obj.test_simple_gather_ndarray() without return types