keras/src/utils/jax_layer.py/JaxLayer/call
class JaxLayer:    def call(self, inputs, training=False):
        def unwrap_variable(variable):
            return None if variable is None else variable.value

        call_args = []
        for argument_name in self.call_fn_arguments:
            if argument_name == "params":
                call_args.append(
                    jax.tree_util.tree_map(unwrap_variable, self.params)
                )
            elif argument_name == "state":
                call_args.append(
                    jax.tree_util.tree_map(unwrap_variable, self.state)
                )
            elif argument_name == "rng":
                call_args.append(self._get_call_rng(training))
            elif argument_name == "inputs":
                call_args.append(inputs)
            elif argument_name == "training":
                call_args.append(training)

        def assign_state_to_variable(value, variable):
            # This exists only to make debugging this error case easier.
            if not hasattr(variable, "assign"):
                raise ValueError(
                    "Structure mismatch: the structure of the state returned "
                    "by `call` does not match the structure of the state at "
                    "initialization time."
                )
            variable.assign(value)

        if self.has_state:
            predictions, new_state = self.call_fn(*call_args)
            jax.tree_util.tree_map(
                assign_state_to_variable, new_state, self.state
            )
            return predictions
        else:
            return self.call_fn(*call_args)