code: def test_optimize_rng_splitting(self):
      def f():
          k1 = base.next_rng_key()
          k2 = base.next_rng_key()
          return k1, k2
  
      key = jax.random.PRNGKey(42)
      assert_allclose = functools.partial(
          np.testing.assert_allclose, atol=1e-5
      )
  
      # With optimize_rng_use the keys returned should be equal to split(n).
      f_opt = transform.transform(random.optimize_rng_use(f))
      jax.tree_util.tree_map(
          assert_allclose,
          f_opt.apply({}, key),
          tuple(jax.random.split(key, 3))[1:],
      )
  
      # Without optimize_rng_use the keys should be equivalent to splitting in a
      # loop.
      f = transform.transform(f)
      jax.tree_util.tree_map(
          assert_allclose, f.apply({}, key), tuple(split_for_n(key, 2))
      )
  
imports: from collections.abc import Sequence
  import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  from haiku._src import base
  from haiku._src import random
  from haiku._src import transform
  import jax
  import jax.extend
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling the RandomTest obj's underlying member method code with example usage: obj.test_optimize_rng_splitting() and returns: k1,  k2