keras/src/ops/core.py/AssociativeScan/compute_output_spec
class AssociativeScan:    def compute_output_spec(self, f, elems, axis):
        elems_flat, tree_ = optree.tree_flatten(elems)
        lens = [elem.shape[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 = optree.tree_unflatten(
            tree_, [slice_along_axis(x, 0, 1, axis=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)
        return y_spec