keras/src/backend/jax/trainer.py/JAXTrainer/compiled_predict_step
class JAXTrainer:    def make_predict_function(self, force=False):
        if self.predict_function is not None and not force:
            return self.predict_function

        def one_predict_step(state, data):
            data = data[0]
            return self.predict_step(state, data)

        def multi_predict_steps(state, data):
            outputs, trainable_variables = one_predict_step(state, data[:1])

            for single_step_data in data[1:]:
                step_outputs, trainable_variables = one_predict_step(
                    state,
                    [single_step_data],
                )
                outputs = tree.map_structure(
                    lambda t1, t2: jax.numpy.concatenate([t1, t2]),
                    outputs,
                    step_outputs,
                )
            return outputs, trainable_variables

        if self.steps_per_execution > 1:
            predict_step = multi_predict_steps
        else:
            predict_step = one_predict_step

        if not self.run_eagerly and self.jit_compile:

            @jax.jit
            def compiled_predict_step(state, data):
                return predict_step(state, data)

            self.predict_function = compiled_predict_step

        else:
            self.predict_function = predict_step