code: def test_key_make_with_custom_impl(self):
      shape = (4, 2, 7)
  
      def seed_rule(_):
          return jnp.ones(shape, dtype=jnp.dtype("uint32"))
  
      def no_rule(*args, **kwargs):
          assert False, "unreachable"
  
      impl = jex.random.define_prng_impl(
          key_shape=shape,
          seed=seed_rule,
          split=no_rule,
          fold_in=no_rule,
          random_bits=no_rule,
      )
      k = jax.random.key(42, impl=impl)
      self.assertEqual(k.shape, ())
      self.assertEqual(impl, jax.random.key_impl(k))
  
file_path: jax/tests/extend_test.py
imports: from absl.testing import absltest
  import jax
  import jax.extend
  import jax.numpy
  from jax._src import abstract_arrays
  from jax._src import linear_util
  from jax._src import prng
  from jax._src import test_util
  from jax import config
label_desc: Use image node_id 1 for calling the RandomTest obj's underlying member method code with example usage: obj.test_key_make_with_custom_impl() and returns: jnp