keras-team/keras · error · ValueError

Expected rebatched data to have batch size 1. Received: shap

Error message

Expected rebatched data to have batch size 1. Received: shape={merged_data.shape}

What it means

When FeatureSpace is used with rebatched=True (e.g. wired as a keras Input via dict_inputs), the internal merge must yield exactly one row so downstream graph construction is unambiguous with output_mode='concat'. A leading batch dimension other than 1 breaks the Functional-model contract.

Source

Thrown at keras/src/layers/preprocessing/feature_space.py:798

            # This scope is to make sure that inner DataLayers
            # will not convert outputs back to backend-native --
            # they should be TF tensors throughout
            preprocessed_data = self._preprocess_features(data)
            preprocessed_data = tree.map_structure(
                lambda x: self._convert_input(x), preprocessed_data
            )

            crossed_data = self._cross_features(preprocessed_data)
            crossed_data = tree.map_structure(
                lambda x: self._convert_input(x), crossed_data
            )

            merged_data = self._merge_features(preprocessed_data, crossed_data)

        if rebatched:
            if self.output_mode == "concat":
                if merged_data.shape[0] != 1:
                    raise ValueError(
                        "Expected rebatched data to have batch size 1. "
                        f"Received: shape={merged_data.shape}"
                    )
                if (
                    backend.backend() != "tensorflow"
                    and not backend_utils.in_tf_graph()
                ):
                    merged_data = np.array(merged_data)
                merged_data = tf.squeeze(merged_data, axis=0)
            else:
                for name, x in merged_data.items():
                    if len(x.shape) == 2 and x.shape[0] == 1:
                        merged_data[name] = tf.squeeze(x, axis=0)

        if (
            backend.backend() != "tensorflow"
            and not backend_utils.in_tf_graph()
        ):

View on GitHub (pinned to 7a34a03db6)

Solutions

  1. Pass data with batch size 1 (a single sample dict) when calling the FeatureSpace directly as an Input
  2. Prefer fs.get_encoded_features() as the model input and feed raw dicts through tf.data instead of calling fs() manually
  3. Check merged_data.shape[0]; unbatch your dataset before the call

Example fix

// before
inputs = fs.get_encoded_features()  # ok
model = keras.Model(inputs, x)
out = model(fs(raw_batched_dict))  # batch > 1 -> error path
// after
out = model(fs(raw_single_row_dict))  # one row at a time
# or build: encoded = fs.get_encoded_features(); model = keras.Model(encoded, x)
Defensive patterns

Strategy: validation

Validate before calling

if isinstance(data, tf.data.Dataset):
    data = data.unbatch().take(1)
assert all(int(v.shape[0]) == 1 for v in batch.values())

Type guard

def is_single_row(d):
    return all(len(v.shape) == 0 or int(v.shape[0]) == 1 for v in d.values())

Try / catch

catch ValueError and pass data with batch size 1 (or the unbatched raw dict for Input usage) instead of a multi-row batch

Prevention

When it happens

Trigger: Passing a batched dataset element (batch_size > 1) or a multi-row dict to a FeatureSpace used as a keras Input, with output_mode='concat'.

Common situations: Using FeatureSpace directly as a keras Input in a Functional model; passing a batched dataset element (batch_size > 1) to that input; rebatching the data with more than one row before the call.

Related errors


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