Skip to content

Preserve single-pass iterables in stack and concat - #105

Open
sylvesterkaczmarek wants to merge 1 commit into
google-research:mainfrom
sylvesterkaczmarek:fix/single-pass-stack-concat
Open

sylvesterkaczmarek wants to merge 1 commit into
google-research:mainfrom
sylvesterkaczmarek:fix/single-pass-stack-concat

Conversation

@sylvesterkaczmarek

Copy link
Copy Markdown

Materialize the input iterable in stack and concat before creating their field-processing callbacks. Currently _ops_base materializes a generator locally, but the callbacks retain the original, exhausted iterator. Both operations then fail while stacking or concatenating the first array field.

The two-line fix shares one materialized list across all array and nested-dataclass fields. List-input behavior, static metadata, axis handling, and backend selection are unchanged.

Validation

macOS arm64 CPU, Python 3.12.11; NumPy 2.5.3, JAX 0.11.2, TensorFlow 2.21.0, PyTorch 2.14.0:

  • All 40 new regression cases fail on unchanged main and pass with the fix.
  • Coverage includes generators, list iterators, a custom single-pass iterable, multiple fields, nested dataclasses, both operations, and two batch axes on all four backends.
  • Full package suite: 425 passed (python -m pytest -q dataclass_array --tb=short).
  • Built and installed the wheel outside the checkout; all 40 new cases pass against the installed package.
  • New-test Pyink formatting, Python compilation, and git diff --check pass. Ruff E9/F passes with F722 excluded for the project's shape-annotation syntax.

Accelerators and documentation builds were not run. This does not change negative-axis handling and is independent of #104.

@sylvesterkaczmarek
sylvesterkaczmarek marked this pull request as ready for review September 23, 2026 08:05
@sylvesterkaczmarek
sylvesterkaczmarek force-pushed the fix/single-pass-stack-concat branch from 79a0f04 to dea42f9 Compare September 30, 2026 14:15
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant