keras/src/utils/jax_layer.py/FlaxLayer/_get_call_rng
class FlaxLayer:    def _get_call_rng(self, training):
        if training:
            return {"dropout": self.seed_generator.next()}
        else:
            return {}