keras-team/keras · error · ValueError

Array inputs to associative_scan must have the same first…

Error message

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

What it means

Consistency guard inside the NumPy `associative_scan`: after flattening the (possibly pytree-structured) `elems` argument to tensors, it checks that every leaf tensor has the same size along the scan axis (axis 0 by default). Because each scan step must combine elements from all leaves, mismatched leading dimensions make the pairing undefined; the message lists the shapes of all leaves it saw.

Solutions

1. Make every tensor in `elems` have the same length along the scan axis (e.g. trim or pad leaves to a common length)
2. Verify the pytree structure of elems: leaves that belong to different structures may be compared unintentionally
3. Set the `axis` argument correctly if the scan axis is not the first dimension

When it happens

Trigger: Thrown at keras/src/backend/numpy/ops/core.py:228 when the library encounters an invalid state.

Common situations: See trigger scenarios.

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

Appendix: source

Thrown at keras/src/backend/numpy/ops/core.py:228

def associative_scan(f, elems, reverse=False, axis=0):
    # Ref: jax.lax.associative_scan
    if not callable(f):
        raise TypeError(f"`f` should be a callable. Received: f={f}")
    elems_flat = tree.flatten(elems)
    elems_flat = [convert_to_tensor(elem) for elem in elems_flat]
    if reverse:
        elems_flat = [np.flip(elem, (axis,)) for elem in elems_flat]

    def _combine(a_flat, b_flat):
        a = tree.pack_sequence_as(elems, a_flat)
        b = tree.pack_sequence_as(elems, b_flat)
        c = f(a, b)
        c_flat = tree.flatten(c)
        return c_flat

    num_elems = int(elems_flat[0].shape[axis])
    if not all(int(elem.shape[axis]) == num_elems for elem in elems_flat[1:]):
        raise ValueError(
            "Array inputs to associative_scan must have the same "
            "first dimension. (saw: {})".format(
                [elem.shape for elem in elems_flat]
            )
        )

    def _interleave(a, b, axis):
        """Given two Tensors of static shape, interleave them along axis."""
        if not (
            a.shape[axis] == b.shape[axis] or a.shape[axis] == b.shape[axis] + 1
        ):
            raise ValueError(
                "Shapes are incompatible for associative_scan interleaving. "
                f"a.shape[{axis}]={a.shape[axis]}, "
                f"b.shape[{axis}]={b.shape[axis]}"
            )

        # we want to get a: [a1, a2], b: [b1, b2]

View on GitHub (pinned to 7a34a03db6)