[Type] Tensor 13: qd.tensor accepts layout= for NDARRAY#533
Closed
hughperkins wants to merge 13 commits intohp/tensor-stork-12from
Closed
[Type] Tensor 13: qd.tensor accepts layout= for NDARRAY#533hughperkins wants to merge 13 commits intohp/tensor-stork-12from
hughperkins wants to merge 13 commits intohp/tensor-stork-12from
Conversation
Removes the PR-6-era NotImplementedError gate. With the AST-rewrite plumbing landed in PR 8 and the cache-key contract pinned in PR 12, the public factory now wires non-identity layouts straight through to ndarray-backed tensors. Behaviour: - shape= is the **canonical** shape the user indexes inside kernels. - The factory allocates the underlying ndarray at the *physical* (permuted) shape: physical[k] = canonical[layout[k]]. - The instance is auto-tagged with _qd_layout, so kernel subscripts x[i, j, ...] are translated to physical access by build_Subscript. - Identity layouts (None or range(ndim)) collapse to no tag — same as the FIELD path — so untagged + identity-tagged stay byte-identical. - order= still forbidden as a kwarg (single source of truth: layout=). - ValueError up front for wrong-length / non-permutation layouts. Tests (test_flexible_factory_layout_ndarray.py, 18 cases): - No-tag / identity collapse - Physical shape and tag presence after rank-2 and rank-3 calls - Validation (length, permutation, order kwarg) - Rank-2 transpose matches the no-layout reference under transpose - Rank-2 explicit value spot-checks - Rank-3 every permutation parametrized - AugAssign through factory-allocated tensors - needs_grad layout inheritance - Cache-key distinction (factory-tagged path) Drive-by: PR-6 NotImplementedError tests in test_flexible_layout.py flipped to assert the factory now succeeds, with depth coverage delegated to the new file. All 180 flexible-tensors tests pass. Doc: flexible_tensors.md "Controlling physical layout" gains an "`layout=` on the ndarray backend" subsection with a worked example and a note about the current physical-shape `tensor.shape` quirk.
Collaborator
Author
|
yay 🙌 |
hughperkins
force-pushed
the
hp/tensor-stork-12
branch
from
e15d272 to
d157522
Compare
hughperkins
force-pushed
the
hp/tensor-stork-13
branch
from
f1b948e to
9f8ce00
Compare
hughperkins
force-pushed
the
hp/tensor-stork-12
branch
from
d157522 to
e15d272
Compare
hughperkins
force-pushed
the
hp/tensor-stork-13
branch
from
9f8ce00 to
f1b948e
Compare
Drops the 'flexible' prefix from filenames and identifiers introduced in this branch series so the user-visible names are simply 'tensor'. Also strips PR-N back-references that will be meaningless once these PRs land. Touches only files owned by this series (no changes to external/ or unrelated tests).
…or-stork-13 Made-with: Cursor # Conflicts: # docs/source/user_guide/tensor.md # python/quadrants/__init__.py # tests/python/test_tensor_layout.py
- Drop the spurious blank line after the test imports so ruff isort is happy. - Assert layout is not None on the non-identity NDARRAY path so pyright can narrow it before the tuple() call (it cannot infer from `order is not None`). Made-with: Cursor
When qd.tensor(..., backend=NDARRAY, layout=..., needs_grad=True) is called, impl.ndarray(..., needs_grad=True) allocates a same-shape companion grad ndarray and wires it via _set_grad. Only the primal was being passed through _with_layout, leaving the grad untagged. As a result, a kernel write to x.grad[i, j, ...] on a non-identity layout would bypass the canonical->physical subscript rewrite and land in the wrong physical slot. Fix: after tagging the primal, propagate the same layout onto arr.grad when present. Tests exercise the tag propagation directly (_qd_layout equality on primal + grad) and via an end-to-end kernel roundtrip on every rank-3 permutation.
… layout grad tests Ndarray.to_numpy() returns the physical (permuted) buffer, as the existing test_factory_layout_rank2_value_check / rank3 tests already demonstrate. The two new ndarray grad tests in test_tensor_layout_grad.py indexed the result with canonical indices, which yielded the wrong slot on rank 2 (assert primal[1, 2] == 12.0 read 21.0) and went out of bounds on most rank-3 permutations. Translate canonical -> physical via the layout permutation when asserting, mirroring the pattern used in test_factory_layout_rank3_all_permutations. Field-backend grad tests keep using canonical indexing (qd.field hides the permutation in to_numpy()). Made-with: Cursor
…ackends test_tensor_layout.py: - Replace test_layout_field_rank3_all_permutations / test_layout_field_rank4_all_permutations with parametrized versions that run on both backends (this only became valid for NDARRAY in the branch that added subscript rewriting for layout-tagged ndarrays; previously it would have raised NotImplementedError). - Replace test_layout_field_with_needs_grad_allocates_grad with a parametrized version covering both backends. - Localise the canonical-vs-physical Ndarray.shape asymmetry inside a small helper so the bug fix later only needs one site updated. test_tensor_layout_grad.py: - Merge the field-only / ndarray-only sister pairs (canonical_kernel_roundtrip_rank2, canonical_kernel_roundtrip_rank3_kij, all_rank3_permutations) into single parametrized tests. - Keep test_layout_ndarray_grad_tag_propagates_rank2 NDARRAY-only with a docstring explaining why: _qd_layout is the NDARRAY mechanism; FIELD uses axis_seq in the SNode tree and the equivalent contract is already covered by the kernel-roundtrip tests above.
Matrix.tensor was added in this stork series; update the expected API list to include it so test_api[arch=x64-Matrix] passes. Made-with: Cursor
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Removes the PR-6-era NotImplementedError gate. With the AST-rewrite plumbing landed in PR 8 and the cache-key contract pinned in PR 12, the public factory now wires non-identity layouts straight through to ndarray-backed tensors.
Behaviour:
Tests (test_flexible_factory_layout_ndarray.py, 18 cases):
Drive-by: PR-6 NotImplementedError tests in test_flexible_layout.py flipped to assert the factory now succeeds, with depth coverage delegated to the new file. All 180 flexible-tensors tests pass.
Doc: flexible_tensors.md "Controlling physical layout" gains an "
layout=on the ndarray backend" subsection with a worked example and a note about the current physical-shapetensor.shapequirk.Issue: #
Brief Summary
copilot:summary
Walkthrough
copilot:walkthrough