keras/src/export/export_lib.py/ExportArchive/track
class ExportArchive:    def track(self, resource):
        """Track the variables (and other assets) of a layer or model.

        By default, all variables used by an endpoint function
        are automatically tracked when you call `add_endpoint()`.
        However, non-variables assets such as lookup tables
        need to be tracked manually. Note that lookup tables
        used by built-in Keras layers
        (`TextVectorization`, `IntegerLookup`, `StringLookup`)
        are automatically tracked in `add_endpoint()`.

        Arguments:
            resource: A trackable TensorFlow resource.
        """
        if backend.backend() == "tensorflow" and not isinstance(
            resource, tf.__internal__.tracking.Trackable
        ):
            raise ValueError(
                "Invalid resource type. Expected an instance of a "
                "TensorFlow `Trackable` (such as a Keras `Layer` or `Model`). "
                f"Received instead an object of type '{type(resource)}'. "
                f"Object received: {resource}"
            )
        if backend.backend() == "jax" and not isinstance(
            resource, backend.jax.layer.JaxLayer
        ):
            raise ValueError(
                "Invalid resource type. Expected an instance of a "
                "JAX-based Keras `Layer` or `Model`. "
                f"Received instead an object of type '{type(resource)}'. "
                f"Object received: {resource}"
            )
        if isinstance(resource, Layer):
            if not resource.built:
                raise ValueError(
                    "The layer provided has not yet been built. "
                    "It must be built before export."
                )

        # Layers in `_tracked` are not part of the trackables that get saved,
        # because we're creating the attribute in a
        # no_automatic_dependency_tracking scope.
        if not hasattr(self, "_tracked"):
            self._tracked = []
        self._tracked.append(resource)

        if isinstance(resource, Layer):
            # Variables in the lists below are actually part of the trackables
            # that get saved, because the lists are created in __init__.
            if backend.backend() == "jax":

                trainable_variables = tree.flatten(resource.trainable_variables)
                non_trainable_variables = tree.flatten(
                    resource.non_trainable_variables
                )
                self._backend_trainable_variables += trainable_variables
                self._backend_non_trainable_variables += non_trainable_variables
                self._backend_variables = (
                    self._backend_trainable_variables
                    + self._backend_non_trainable_variables
                )

                self._tf_trackable.trainable_variables += [
                    tf.Variable(v) for v in trainable_variables
                ]
                self._tf_trackable.non_trainable_variables += [
                    tf.Variable(v) for v in non_trainable_variables
                ]
                self._tf_trackable.variables = (
                    self._tf_trackable.trainable_variables
                    + self._tf_trackable.non_trainable_variables
                )
            else:
                self._tf_trackable.variables += resource.variables
                self._tf_trackable.trainable_variables += (
                    resource.trainable_variables
                )
                self._tf_trackable.non_trainable_variables += (
                    resource.non_trainable_variables
                )