keras/src/backend/jax/trainer.py/JAXTrainer/__init__
class JAXTrainer:    def __init__(self):
        super().__init__()
        self.train_function = None
        self.test_function = None
        self.predict_function = None
        self._jax_state_synced = True