keras-team/keras · error · NotImplementedError

The TFSMLayer is only currently supported with the TensorFlo

Error message

The TFSMLayer is only currently supported with the TensorFlow backend.

What it means

TFSMLayer reloads a TensorFlow SavedModel, which requires the TensorFlow runtime, so keras.src.export.tfsm_layer.__init__ raises NotImplementedError when keras.backend() is anything other than 'tensorflow'. Keras 3 is multi-backend, and the reloader only works when the active backend is TensorFlow. The check happens before any loading, so it fires immediately at construction.

Source

Thrown at keras/src/export/tfsm_layer.py:59

    custom signature.
    * If you need training-time behavior to differ from inference-time behavior
    (i.e. if you need the reloaded object to support a `training=True` argument
    in `__call__()`), make sure that the training-time call function is
    saved as a standalone endpoint in the artifact, and provide its name
    to the `TFSMLayer` via the `call_training_endpoint` argument.
    """

    def __init__(
        self,
        filepath,
        call_endpoint="serve",
        call_training_endpoint=None,
        trainable=True,
        name=None,
        dtype=None,
    ):
        if backend.backend() != "tensorflow":
            raise NotImplementedError(
                "The TFSMLayer is only currently supported with the "
                "TensorFlow backend."
            )

        # Initialize an empty layer, then add_weight() etc. as needed.
        super().__init__(trainable=trainable, name=name, dtype=dtype)

        self._reloaded_obj = tf.saved_model.load(filepath)

        self.filepath = filepath
        self.call_endpoint = call_endpoint
        self.call_training_endpoint = call_training_endpoint

        # Resolve the call function.
        if hasattr(self._reloaded_obj, call_endpoint):
            # Case 1: it's set as an attribute.
            self.call_endpoint_fn = getattr(self._reloaded_obj, call_endpoint)
        elif call_endpoint in self._reloaded_obj.signatures:

View on GitHub (pinned to 7a34a03db6)

Solutions

  1. Set the backend before any keras import: os.environ['KERAS_BACKEND']='tensorflow' at the top of the entry script, then restart/reimport.
  2. Convert the SavedModel to a backend-neutral format instead: reload under TF once, save as .keras (model.save), then load in your jax or torch pipeline.
  3. Verify with keras.config.backend() right before constructing TFSMLayer to catch backend drift early.

Example fix

# before (KERAS_BACKEND=jax was set)
layer = keras.layers.TFSMLayer('saved_model/')  # -> NotImplementedError

# after
import os
os.environ['KERAS_BACKEND'] = 'tensorflow'
import keras
layer = keras.layers.TFSMLayer('saved_model/')
Defensive patterns

Strategy: validation

Validate before calling

import keras
assert keras.config.backend() == 'tensorflow', 'TFSMLayer requires the TF backend'

Prevention

When it happens

Trigger: Constructing keras.layers.TFSMLayer('path/to/saved_model', ...) while KERAS_BACKEND is 'jax', 'torch', or 'numpy'; common when a notebook set the backend env var earlier or the script imports a jax-based library first.

Common situations: Teams migrating to Keras 3 with KERAS_BACKEND=jax or torch in CI; loading a TF-Hub or legacy TF2 SavedModel in a JAX training pipeline; default backend resolution picking a non-TF backend installed alongside TF.

Related errors


AI-assisted analysis of keras-team/keras@7a34a03db6 (2026-08-25). Data as JSON: /api/errors/5707e3060100aecc. Report an issue: GitHub.