keras/src/utils/jax_layer.py/JaxLayer/_get_init_rng
class JaxLayer:    def _get_init_rng(self):
        """
        Returns a JAX `PRNGKey` or structure of `PRNGKey`s to pass to `init_fn`.

        By default, this returns a single `PRNGKey` retrieved by calling
        `self.seed_generator.next()`. Override this to return a different
        structure.

        Returns:
            a JAX `PRNGKey` or structure of `PRNGKey`s that will be passed as
            the `rng` argument of `init_fn`.
        """
        return self.seed_generator.next()