Keras
Machine Learning
Neural Networks
Optional Inputs
Python

How to create Keras model with optional inputs

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Keras models do not treat declared Input tensors as truly optional. Once a functional model is built with two inputs, both inputs are part of the signature. The usual workaround is to always provide the auxiliary input with a default value and a mask, or to design separate model paths for the "with extra input" and "without extra input" cases.

Why Functional Inputs Are Not Optional

In the functional API, the input signature is fixed when you build the graph.

python
1import tensorflow as tf
2
3main_input = tf.keras.Input(shape=(8,), name="main")
4aux_input = tf.keras.Input(shape=(4,), name="aux")

If you create a model from these two inputs, Keras expects both every time. You cannot omit aux during fit() or predict() and expect the graph to rewire itself dynamically.

That is why the real problem is usually data representation, not optional function arguments in the Python sense.

Pattern 1: Always Pass the Auxiliary Tensor and a Mask

The most common solution is to always pass an auxiliary tensor. When the extra data is missing, pass zeros and provide a mask that says whether the values are real.

python
1import tensorflow as tf
2
3main_input = tf.keras.Input(shape=(8,), name="main")
4aux_input = tf.keras.Input(shape=(4,), name="aux")
5aux_present = tf.keras.Input(shape=(1,), name="aux_present")
6
7x = tf.keras.layers.Concatenate()([main_input, aux_input, aux_present])
8x = tf.keras.layers.Dense(16, activation="relu")(x)
9output = tf.keras.layers.Dense(1, activation="sigmoid")(x)
10
11model = tf.keras.Model(
12    inputs=[main_input, aux_input, aux_present],
13    outputs=output,
14)
15
16model.compile(optimizer="adam", loss="binary_crossentropy")

When the auxiliary input is missing:

python
1prediction = model.predict(
2    {
3        "main": tf.random.normal((2, 8)),
4        "aux": tf.zeros((2, 4)),
5        "aux_present": tf.zeros((2, 1)),
6    },
7    verbose=0,
8)

This avoids signature mismatch and lets the model learn the difference between "missing" and "present."

Pattern 2: Use a Subclassed Model with a Default

If you need Python-level flexibility, a subclassed model can inject a default tensor when an auxiliary input is absent.

python
1import tensorflow as tf
2
3
4class OptionalAuxModel(tf.keras.Model):
5    def __init__(self):
6        super().__init__()
7        self.hidden = tf.keras.layers.Dense(16, activation="relu")
8        self.out = tf.keras.layers.Dense(1, activation="sigmoid")
9
10    def call(self, inputs):
11        main = inputs["main"]
12        aux = inputs.get("aux")
13
14        if aux is None:
15            aux = tf.zeros((tf.shape(main)[0], 4), dtype=main.dtype)
16
17        x = tf.concat([main, aux], axis=-1)
18        x = self.hidden(x)
19        return self.out(x)
20
21
22model = OptionalAuxModel()
23result = model({"main": tf.random.normal((2, 8))})
24print(result)

This is flexible, but it is more manual than a standard functional graph and can complicate serialization or tooling expectations if overused.

Pattern 3: Separate Models

Sometimes the cleanest design is two models:

  • one model that expects only the main input
  • one model that expects main plus auxiliary input

This is often easier when the "with aux" and "without aux" cases are semantically different enough that sharing one graph becomes awkward.

The right choice depends on whether the auxiliary input is truly optional metadata or whether it changes the problem definition itself.

Keep the Data Pipeline Consistent

The model architecture is only part of the solution. Your training pipeline also has to represent missing auxiliary data consistently. If some batches contain one structure and others contain another, training code becomes fragile quickly.

In practice, most teams normalize the pipeline so every example has the same fields, even if one field sometimes contains defaults plus a presence mask.

Common Pitfalls

Declaring an auxiliary input in the functional API and then omitting it entirely at call time causes a signature mismatch because Keras still expects it.

Using all-zero auxiliary data without a mask is ambiguous if zero is also a valid real value for that feature.

Building a highly flexible subclassed model when two separate simpler models would be clearer can make training and debugging harder.

Ignoring how the training data represents missing values often creates mismatch between training-time behavior and inference-time behavior.

Thinking of Keras inputs as optional Python parameters rather than fixed graph inputs leads to the wrong mental model.

Summary

  • Functional Keras inputs are fixed once the model graph is built.
  • The usual workaround is to always provide the input and represent absence with defaults plus a mask.
  • A subclassed model can inject defaults dynamically if you need more call-time flexibility.
  • Separate models may be cleaner when the optional path really represents a different problem.
  • Keep the data pipeline consistent so missing auxiliary data is represented the same way during training and inference.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.