keras-team/keras · error · ValueError

Array inputs to associative_scan must have the same first di

Error message

Array inputs to associative_scan must have the same first dimension. (saw: {})

What it means

associative_scan folds a function over parallel arrays and requires every element in the (possibly nested) elems structure to share the same leading dimension (the scan length). Mismatched leading dims make the scan ill-defined.

Source

Thrown at keras/src/ops/core.py:210

    )


class AssociativeScan(Operation):
    def __init__(self, reverse=False, axis=0, *, name=None):
        super().__init__(name=name)
        self.reverse = reverse
        self.axis = axis

    def call(self, f, elems):
        return backend.core.associative_scan(
            f, elems, reverse=self.reverse, axis=self.axis
        )

    def compute_output_spec(self, f, elems):
        elems_flat = tree.flatten(elems)
        lens = [elem.shape[self.axis] for elem in elems_flat]
        if len(set(lens)) != 1:
            raise ValueError(
                "Array inputs to associative_scan must have the same "
                "first dimension. (saw: {})".format(
                    [elem.shape for elem in elems_flat]
                )
            )

        x = tree.pack_sequence_as(
            elems,
            [slice_along_axis(x, 0, 1, axis=self.axis) for x in elems_flat],
        )
        y_spec = backend.compute_output_spec(f, x, x)

        def _restore_shape(x):
            return KerasTensor(
                shape=elems_flat[0].shape, dtype=x.dtype, sparse=x.sparse
            )

        y_spec = tree.map_structure(_restore_shape, y_spec)

View on GitHub (pinned to 7a34a03db6)

Solutions

1. Slice or pad the arrays so all leading dims match
2. Fix the data pipeline that produced inconsistent lengths
3. If using tuples of (elems, init), verify init is not accidentally included in the scan structure

Example fix

# before
keras.ops.associative_scan(fn, (xs, ys))  # len(xs) != len(ys)

# after
n = min(len(xs), len(ys))
keras.ops.associative_scan(fn, (xs[:n], ys[:n]))

lens = {e.shape[0] for e in tree.flatten(elems)}
assert len(lens) == 1, f'mismatched leading dims: {lens}'

When it happens

Trigger: keras.ops.associative_scan(f, (a, b)) where a.shape[0] != b.shape[0]

Common situations: JAX-style functional loops ported to Keras 3 ops, sequence models with padded batches

Related errors

• The number of dimensions in `inputs` must match the number o
• Layer {self.name} weight shape {variable.shape} is not compa
• Expected rebatched data to have batch size 1. Received: shap
• Expected as input a list/tuple of 2 tensors. Received input_
• Expected the two input tensors to have identical shapes. Rec

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