keras/src/utils/jax_layer.py/JaxLayer/build
class JaxLayer:    def build(self, input_shape):
        if self.params is not None or self.state is not None:
            return

        if jax_utils.is_in_jax_tracing_scope():
            # This exception is not actually shown, it is caught and a detailed
            # warning about calling 'build' is printed.
            raise ValueError("'JaxLayer' cannot be built in tracing scope")

        # Initialize `params` and `state` if needed by calling `init_fn`.
        def create_input(shape):
            shape = [d if d is not None else 1 for d in shape]
            return jax.numpy.ones(shape)

        init_inputs = tree.map_shape_structure(create_input, input_shape)
        init_args = []
        for argument_name in self.init_fn_arguments:
            if argument_name == "rng":
                init_args.append(self._get_init_rng())
            elif argument_name == "inputs":
                init_args.append(init_inputs)
            elif argument_name == "training":
                init_args.append(True)

        init_result = self.init_fn(*init_args)
        if self.has_state:
            init_params, init_state = init_result
        else:
            init_params, init_state = init_result, None

        self.tracked_params = self._create_variables(
            init_params, trainable=True
        )
        self.tracked_state = self._create_variables(init_state, trainable=False)
        self.built = True