keras/src/utils/jax_layer.py/JaxLayer/__init__
class JaxLayer:    def __init__(
        self,
        call_fn,
        init_fn=None,
        params=None,
        state=None,
        seed=None,
        **kwargs,
    ):
        if backend.backend() != "jax":
            raise ValueError(
                "JaxLayer is only supported with the JAX backend. Current "
                f"backend: {backend.backend()}"
            )

        if init_fn is None and params is None and state is None:
            raise ValueError(
                "`init_fn`, `params` and `state` cannot all be `None`."
            )

        super().__init__(**kwargs)
        self.call_fn = call_fn
        self.init_fn = init_fn
        self.seed_generator = backend.random.SeedGenerator(seed)
        self.tracked_params = self._create_variables(params, trainable=True)
        self.tracked_state = self._create_variables(state, trainable=False)
        if self.params is not None or self.state is not None:
            self.built = True

        self.call_fn_arguments = self._validate_signature(
            call_fn,
            "call_fn",
            {"params", "state", "rng", "inputs", "training"},
            {"inputs"},
        )
        self.has_state = "state" in self.call_fn_arguments

        if init_fn:
            self.init_fn_arguments = self._validate_signature(
                init_fn, "init_fn", {"rng", "inputs", "training"}, {"inputs"}
            )