skip to content

Why does calling framework ops inside a Keras 3 custom layer break portability?

level: seniorimportance: should knowfreq 42%

answer

  1. one API, three array types underneath
  2. a NumPy-shaped namespace inside Keras
  3. random numbers have their own namespace
  4. loads fine, fails on the first forward pass

basics

~20 s

A layer whose call() invokes one framework's tensor API only runs when that framework is the active backend, so the saved model silently stops being portable. keras.ops is the backend-agnostic numerics layer, and keras.random is its counterpart for random numbers.

solid answer

~40 s

Keras 3 dispatches to whichever backend the process selected, so the tensors flowing through `call()` are that backend's arrays. Reach past Keras to TensorFlow's, JAX's or PyTorch's own ops and you have hard-wired the layer: under any other backend those functions receive an array type they cannot consume and the layer breaks. `keras.ops` exists to prevent this - a NumPy-shaped API (`keras.ops.matmul`, `keras.ops.mean`, `keras.ops.softmax`, `keras.ops.convert_to_tensor`) that lowers to the active backend. Randomness has its own namespace, `keras.random`, with a `SeedGenerator` so behaviour is reproducible rather than depending on a framework's global RNG. The archive stays loadable either way, which is the trap: portability breaks at the first forward pass, not at load.

code

python · 19 lines
python
import keras

@keras.saving.register_keras_serializable(package="MyPkg")
class NoisyScale(keras.layers.Layer):
    def __init__(self, factor=1.0, **kwargs):
        super().__init__(**kwargs)
        self.factor = factor
        self.seed_generator = keras.random.SeedGenerator(1337)

    def call(self, x, training=False):
        y = keras.ops.multiply(x, self.factor)
        if training:
            y = y + keras.random.normal(keras.ops.shape(y), seed=self.seed_generator)
        return y

    def get_config(self):
        config = super().get_config()
        config["factor"] = self.factor
        return config

go deeper

for a junior

Know that custom layer maths should use keras.ops rather than a specific framework's tensor functions, so the same layer runs on any Keras 3 backend.

for a middle

Explain the mechanism: tensors in call() belong to the active backend, so a foreign op receives a type it cannot handle, and keras.ops lowers to whichever backend is live.

for a senior

Point out that the failure surfaces at the first forward pass rather than at load, and prescribe the guard: keras.random for RNG plus CI that exercises custom layers under a second backend.

for a principal

Treat backend independence as an option worth paying for. Set the standard that shared layer code uses keras.ops, and require any deliberate pin to be documented and checked so the organisation can still change backends later.

## Where portability is actually decided A `.keras` archive is backend-neutral by construction: class names, arguments, numeric arrays. So people conclude their models are portable. They are - right up to the first custom layer whose call() reaches for a specific framework. Under Keras 3, whatever enters call() is the active backend's array type. Passing that to a different framework's op fails: the function does not recognise the type, and you get a type error out of the depths of a library you did not think you were using. The layer's config loaded perfectly; the maths cannot run. ## keras.ops `keras.ops` is the backend-agnostic numerics layer, deliberately shaped like NumPy: `keras.ops.matmul`, `keras.ops.mean`, `keras.ops.reshape`, `keras.ops.concatenate`, `keras.ops.where`, plus neural-network ops such as `keras.ops.softmax` and `keras.ops.relu`. Each call lowers to the active backend's implementation. `keras.ops.convert_to_tensor` and `keras.ops.convert_to_numpy` are the entry and exit points when you must interoperate with plain arrays. Write call() against keras.ops and the layer is portable by construction - the same source runs under all three backends, and the saved model loads and executes anywhere. ## keras.random Randomness deserves its own mention because reaching for a framework's global RNG is both a portability bug and a reproducibility bug. `keras.random` provides normal, uniform, dropout and friends, and `keras.random.SeedGenerator` gives a layer seed state it owns. That matters beyond portability: JAX's functional model has no global mutable RNG to lean on, so a layer written against one framework's stateful generator has no faithful equivalent there. Creating a SeedGenerator in `__init__` and drawing from it in call() is the pattern that behaves the same everywhere. ## When it is acceptable to pin Sometimes the escape hatch is the right answer: an op with no keras.ops equivalent, a custom kernel, a framework-specific primitive your model genuinely needs. Then be explicit. Say in the class docstring that this layer requires that backend, check `keras.backend.backend()` and raise a clear error rather than letting users discover it as a type error, and keep the pinned code in one narrow layer instead of scattering it. The cost is real - the model can no longer be handed to a team on a different stack - so it should be a decision, not an accident. ## The Keras 2 habit Most violations are muscle memory. In Keras 2, Keras was TensorFlow's front end and mixing in TensorFlow ops was normal, even idiomatic. That code keeps working under the TensorFlow backend, so nothing forces the issue until someone tries another backend or inherits the model. A CI job that runs the custom-layer unit tests under a second backend turns this from an archaeology exercise into a failing test on the day it is written. ## How the failure presents Worth rehearsing for an interview: it does not present as a saving problem. load_model succeeds, summary() prints, and the process dies on the first predict() with a type error naming a framework nobody expected to be in the stack. Recognising that shape of failure - loads fine, runs wrong - is the senior signal here.

  • Why does randomness need keras.random rather than a framework's global RNG?
    Because backends disagree about RNG state. JAX is functional and threads keys explicitly, while others keep a mutable global generator, so a layer written against one has no faithful equivalent on the other. keras.random with a SeedGenerator gives the layer its own seed state, which is both portable and reproducible across runs and backends.
  • How do you make a deliberately backend-pinned layer safe to ship?
    Make the dependency loud. Document it on the class, check keras.backend.backend() at construction and raise a clear error naming the required backend, and isolate the pinned code in one layer rather than spreading it. A user then hits an explanatory message instead of a type error from a library they never imported.
  • Would this problem be caught by loading the model?
    No. The archive stores a class name and a config, so load_model resolves the class and rebuilds the layer without complaint. The framework-specific call lives in call(), so the failure appears on the first forward pass. Catch it by running the layer's unit tests under a second backend in CI.

saying these in an interview costs you the question

  • Assuming tensors in a Keras 3 layer are always TensorFlow tensors
  • Believing Keras auto-converts foreign framework ops
  • Using a framework's global RNG inside call()
  • Expecting load_model to surface a backend-pinned layer

context

resolving links…