keras/src/utils/jax_layer.py/FlaxLayer/_get_init_rng
class FlaxLayer:    def _get_init_rng(self):
        return {
            "params": self.seed_generator.next(),
            "dropout": self.seed_generator.next(),
        }