keras/src/export/export_lib.py/ExportArchive/add_endpoint
class ExportArchive:    def add_endpoint(self, name, fn, input_signature=None, jax2tf_kwargs=None):
        """Register a new serving endpoint.

        Arguments:
            name: Str, name of the endpoint.
            fn: A function. It should only leverage resources
                (e.g. `tf.Variable` objects or `tf.lookup.StaticHashTable`
                objects) that are available on the models/layers
                tracked by the `ExportArchive` (you can call `.track(model)`
                to track a new model).
                The shape and dtype of the inputs to the function must be
                known. For that purpose, you can either 1) make sure that
                `fn` is a `tf.function` that has been called at least once, or
                2) provide an `input_signature` argument that specifies the
                shape and dtype of the inputs (see below).
            input_signature: Used to specify the shape and dtype of the
                inputs to `fn`. List of `tf.TensorSpec` objects (one
                per positional input argument of `fn`). Nested arguments are
                allowed (see below for an example showing a Functional model
                with 2 input arguments).
            jax2tf_kwargs: Optional. A dict for arguments to pass to `jax2tf`.
                Supported only when the backend is JAX. See documentation for
                [`jax2tf.convert`](
                    https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md).
                The values for `native_serialization` and `polymorphic_shapes`,
                if not provided, are automatically computed.

        Returns:
            The `tf.function` wrapping `fn` that was added to the archive.

        Example:

        Adding an endpoint using the `input_signature` argument when the
        model has a single input argument:

        ```python
        export_archive = ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint(
            name="serve",
            fn=model.call,
            input_signature=[tf.TensorSpec(shape=(None, 3), dtype=tf.float32)],
        )
        ```

        Adding an endpoint using the `input_signature` argument when the
        model has two positional input arguments:

        ```python
        export_archive = ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint(
            name="serve",
            fn=model.call,
            input_signature=[
                tf.TensorSpec(shape=(None, 3), dtype=tf.float32),
                tf.TensorSpec(shape=(None, 4), dtype=tf.float32),
            ],
        )
        ```

        Adding an endpoint using the `input_signature` argument when the
        model has one input argument that is a list of 2 tensors (e.g.
        a Functional model with 2 inputs):

        ```python
        model = keras.Model(inputs=[x1, x2], outputs=outputs)

        export_archive = ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint(
            name="serve",
            fn=model.call,
            input_signature=[
                [
                    tf.TensorSpec(shape=(None, 3), dtype=tf.float32),
                    tf.TensorSpec(shape=(None, 4), dtype=tf.float32),
                ],
            ],
        )
        ```

        This also works with dictionary inputs:

        ```python
        model = keras.Model(inputs={"x1": x1, "x2": x2}, outputs=outputs)

        export_archive = ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint(
            name="serve",
            fn=model.call,
            input_signature=[
                {
                    "x1": tf.TensorSpec(shape=(None, 3), dtype=tf.float32),
                    "x2": tf.TensorSpec(shape=(None, 4), dtype=tf.float32),
                },
            ],
        )
        ```

        Adding an endpoint that is a `tf.function`:

        ```python
        @tf.function()
        def serving_fn(x):
            return model(x)

        # The function must be traced, i.e. it must be called at least once.
        serving_fn(tf.random.normal(shape=(2, 3)))

        export_archive = ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint(name="serve", fn=serving_fn)
        ```
        """
        if name in self._endpoint_names:
            raise ValueError(f"Endpoint name '{name}' is already taken.")

        if jax2tf_kwargs and backend.backend() != "jax":
            raise ValueError(
                "'jax2tf_kwargs' is only supported with the jax backend. "
                f"Current backend: {backend.backend()}"
            )

        if input_signature:
            if backend.backend() == "tensorflow":
                decorated_fn = tf.function(
                    fn, input_signature=input_signature, autograph=False
                )
            else:  # JAX backend

                # 1. Create a stateless wrapper for `fn`
                # 2. jax2tf the stateless wrapper
                # 3. Create a stateful function that binds the variables with
                #    the jax2tf converted stateless wrapper
                # 4. Make the signature of the stateful function the same as the
                #    original function
                # 5. Wrap in a `tf.function`
                def stateless_fn(variables, *args, **kwargs):
                    state_mapping = zip(self._backend_variables, variables)
                    with StatelessScope(state_mapping=state_mapping) as scope:
                        output = fn(*args, **kwargs)

                    # Gather updated non-trainable variables
                    non_trainable_variables = []
                    for var in self._backend_non_trainable_variables:
                        new_value = scope.get_current_value(var)
                        non_trainable_variables.append(new_value)
                    return output, non_trainable_variables

                jax2tf_stateless_fn = self._convert_jax2tf_function(
                    stateless_fn,
                    input_signature,
                    jax2tf_kwargs=jax2tf_kwargs,
                )

                def stateful_fn(*args, **kwargs):
                    output, non_trainable_variables = jax2tf_stateless_fn(
                        # Change the trackable `ListWrapper` to a plain `list`
                        list(self._tf_trackable.variables),
                        *args,
                        **kwargs,
                    )
                    for var, new_value in zip(
                        self._tf_trackable.non_trainable_variables,
                        non_trainable_variables,
                    ):
                        var.assign(new_value)
                    return output

                # Note: we truncate the number of parameters to what is
                # specified by `input_signature`.
                fn_signature = inspect.signature(fn)
                fn_parameters = list(fn_signature.parameters.values())
                stateful_fn.__signature__ = inspect.Signature(
                    parameters=fn_parameters[0 : len(input_signature)],
                    return_annotation=fn_signature.return_annotation,
                )

                decorated_fn = tf.function(
                    stateful_fn,
                    input_signature=input_signature,
                    autograph=False,
                )
            self._endpoint_signatures[name] = input_signature
        else:
            if isinstance(fn, tf.types.experimental.GenericFunction):
                if not fn._list_all_concrete_functions():
                    raise ValueError(
                        f"The provided tf.function '{fn}' "
                        "has never been called. "
                        "To specify the expected shape and dtype "
                        "of the function's arguments, "
                        "you must either provide a function that "
                        "has been called at least once, or alternatively pass "
                        "an `input_signature` argument in `add_endpoint()`."
                    )
                decorated_fn = fn
            else:
                raise ValueError(
                    "If the `fn` argument provided is not a `tf.function`, "
                    "you must provide an `input_signature` argument to "
                    "specify the shape and dtype of the function arguments. "
                    "Example:\n\n"
                    "export_archive.add_endpoint(\n"
                    "    name='call',\n"
                    "    fn=model.call,\n"
                    "    input_signature=[\n"
                    "        tf.TensorSpec(\n"
                    "            shape=(None, 224, 224, 3),\n"
                    "            dtype=tf.float32,\n"
                    "        )\n"
                    "    ],\n"
                    ")"
                )
        setattr(self._tf_trackable, name, decorated_fn)
        self._endpoint_names.append(name)
        return decorated_fn