code: def _test_repeat(self, x1, x2, layer, n, rng_params, **kwargs):
      init_fn, apply_fn, kernel_fn = (
          stax.Identity() if n == 0 else stax.serial(*([layer] * n))
      )
      init_fn_repeat, apply_fn_repeat, kernel_fn_repeat = stax.repeat(
          layer, n
      )
  
      out_shape, params = init_fn(rng_params, x1.shape)
      out_shape_repeat, params_repeat = init_fn_repeat(
          rng_params, x1.shape
      )
  
      self.assertEqual(out_shape, out_shape_repeat)
  
      kwargs1 = {k: kwargs[k][0] for k in kwargs}
      out = apply_fn(params, x1, **kwargs1)
      out_repeat = apply_fn_repeat(params_repeat, x1, **kwargs1)
  
      self.assertAllClose(out, out_repeat)
  
      for get in [
          None,
          "ntk",
          "nngp",
          "cov1",
          ("nngp", "cov1"),
          ("cov1", "ntk"),
      ]:
          with self.subTest(get=get):
              k = kernel_fn(x1, x2, get, **kwargs)
              k_repeat = kernel_fn_repeat(x1, x2, get, **kwargs)
              self.assertAllClose(k, k_repeat)
  
file_path: neural-tangents/tests/stax/combinators_test.py
imports: import random
  from absl.testing import absltest
  from jax import config
  from jax import random
  import jax.numpy
  from neural_tangents import stax
  from tests import test_utils
label_desc: Use image node_id 1 for calling the RepeatTest obj's underlying member method code with example usage: obj._test_repeat(x1, x2, layer, n, rng_params) without return types