[Fix][Reduction] align constructor dim defaults + dim=[] noop on All/Any#1435
Conversation
lcy-seso
added
automated
|
/gemini review |
There was a problem hiding this comment.
Pull request overview
Aligns the Python reduction-op implementations with the manifest’s declared interface by (1) making dim=None the default for ops whose manifest declares default: null (full reduction), while preserving ProdFwdOp(dim=-1), and (2) implementing the manifest/PyTorch no-op semantics for AllFwdOp / AnyFwdOp when dim=[] / dim=().
Changes:
- Flips constructor defaults
dim=-1 → Nonefor the affected reduce-op bases and logical/count reduction ops; adds an explicitProdFwdOp(dim=-1)override. - Extends
EmptyDimPolicywith"noop"and adds an op-layer no-op short-circuit for empty-dim reductions (casting output dtype for All/Any). - Updates/extends tests to pin default-dim and empty-dim semantics and to make prior
dim=-1expectations explicit at call sites.
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tileops/ops/reduction/reduce.py | Changes default dim to None, adds empty-dim no-op short-circuit, and preserves ProdFwdOp(dim=-1) via an explicit __init__. |
| tileops/ops/reduction/_multidim.py | Extends EmptyDimPolicy with "noop" and updates normalize_dim() behavior for empty dims under that policy. |
| tileops/ops/reduction/all_op.py | Sets _empty_dim_policy="noop", updates default dim=None, and ensures noop path outputs bool. |
| tileops/ops/reduction/any_op.py | Sets _empty_dim_policy="noop", updates default dim=None, and ensures noop path outputs bool. |
| tileops/ops/reduction/count_nonzero.py | Updates default dim=None (full reduction) while retaining keepdim=False behavior. |
| tests/ops/test_reduction_defaults.py | Adds regression coverage for default dim semantics and All/Any empty-dim no-op behavior. |
| tests/ops/test_reduce.py | Updates tests to pass dim=-1 explicitly where that behavior is intended. |
| tests/ops/test_logical_reduce.py | Updates logical-reduce tests to pass dim=-1 explicitly where that behavior is intended. |
| tests/ops/test_reduce_multidim.py | Updates the All empty-dim test to assert the new no-op semantics. |
Comments suppressed due to low confidence (1)
tileops/ops/reduction/_multidim.py:32
normalize_dimis documented/used as acceptingdim=[]anddim=()(tuple), but the type annotation isUnion[int, list[int], None]. This is now inconsistent with both the docstring and call sites that pass tuples; please widen the annotation (e.g., to includetuple[int, ...]or aSequence[int]) so the public helper matches the supported API surface.
def normalize_dim(
dim: Union[int, list[int], None],
ndim: int,
*,
empty_dim_policy: EmptyDimPolicy = "reject",
) -> list[int]:
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request updates the default behavior of reduction operations to perform a full reduction (dim=None) by default, aligning with standard PyTorch semantics, while ProdFwdOp retains its dim=-1 default. It introduces a no-op contract for AllFwdOp and AnyFwdOp when an empty dimension list is provided, allowing them to return the input tensor cast to boolean. The implementation includes a new short-circuit mechanism in _ReduceOpBase and updates to dimension normalization logic. A comprehensive set of regression tests has been added to verify these constructor defaults and empty-dimension behaviors. Feedback was provided to refactor duplicated input validation logic into a shared private method to improve code maintainability.
There was a problem hiding this comment.
Code Review
This pull request updates the default reduction dimension from -1 to None (full reduction) for the majority of reduction operations, with ProdFwdOp being the exception. It also introduces a 'noop' contract for AllFwdOp and AnyFwdOp when dim=[] is passed, returning the input cast to boolean. Feedback points out that the change in default dimension is a breaking change for external consumers and suggests that the roofline model for no-op reductions should be adjusted to better account for input memory bandwidth.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 10 out of 10 changed files in this pull request and generated 3 comments.
Comments suppressed due to low confidence (1)
tileops/ops/reduction/_multidim.py:32
normalize_dim()now explicitly documents/supportsdim=()and callers pass tuples, but the function’s type annotation is stillUnion[int, list[int], None]. This makes the public typing inconsistent with actual behavior and the manifest interface; please widen it to includetuple[int, ...](orSequence[int]).
def normalize_dim(
dim: Union[int, list[int], None],
ndim: int,
*,
empty_dim_policy: EmptyDimPolicy = "reject",
) -> list[int]:
Aligns the ten reduction ops whose manifest declares ``default: null`` on ``dim`` with the manifest spec by switching their constructor default from ``-1`` to ``None`` (full reduction). ProdFwdOp keeps its documented ``dim=-1`` default via an explicit subclass override. AllFwdOp / AnyFwdOp opt into a new ``EmptyDimPolicy="noop"`` branch so that ``dim=[]`` / ``dim=()`` returns the input cast to bool, matching PyTorch's identity semantics. normalize_dim() learns the third policy value; _ReduceOpBase grows a _maybe_noop short-circuit in forward(). Test callsites that omitted ``dim=`` and relied on the old ``-1`` default are updated to pass ``dim=-1`` explicitly, preserving test intent.
The empty-dim noop short-circuit on AllFwdOp / AnyFwdOp was returning the input directly without running the standard CUDA/dtype/ndim validation and without binding _last_roofline_mn. As a result a cpu tensor or wrong-dtype tensor with dim=[] silently passed through, and eval_roofline() after a noop forward raised RuntimeError because the shape state was never populated. _maybe_noop now mirrors the validation that _prepare_input runs, and binds _last_roofline_mn = (numel, 0) so the public forward contract is preserved (bad inputs raise; eval_roofline() works after a noop). Adds six regression tests covering both ops: cpu-tensor rejection, wrong-dtype rejection, and roofline-binding after a successful noop forward.
… tests - Widen `dim` type annotation to include `tuple[int, ...]` in `_ReduceOpBase`, `ProdFwdOp`, `_WelfordReduceOp`, `AllFwdOp`, `AnyFwdOp`, `CountNonzeroFwdOp` (runtime path already accepts tuples via normalize_dim; annotation now matches manifest tuple[int, ...] token). - Pass `dim=-1` explicitly at the six remaining last-axis callsites in `tests/ops/test_welford_non_aligned.py` (lines 204/219/234/249/262/275), matching the audit-list policy and the file's own "single-dim, dim=-1" header comment. - Extract `_validate_input_tensor(x)` helper shared by `_prepare_input` and `_maybe_noop`, removing duplicated CUDA/dtype/ndim checks. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
…-zero The dim=[] noop short-circuit previously bound _last_roofline_mn = (numel, 0). Under the per-op-kind formulas in eval_roofline(), N=0 collapses both flops and the M*N input-read term to zero, under-counting the actual data movement: the noop still reads every input element and writes an equal-shape result (cast to bool for All/Any/CountNonzero, identity for Sum-style ops). Bind (numel, 1) instead: model the noop as a degenerate reduction over an axis of length 1. Per the existing formulas this yields mem_bytes ~ numel * elem_bytes + numel (read + output term) rather than zero, putting the noop on the correct bandwidth scale. Update test_all_empty_dim_noop_binds_roofline / test_any_empty_dim_noop_binds_roofline to assert mem_bytes is at least the input-read term and at most 2x read + output write -- catches future regressions to N=0.
Surface the BREAKING CHANGE trailer so release-notes tooling (conventional-commits / semantic-release) picks up the constructor- default flip for the ten reduce ops whose manifest declares `default: null`. The original implementation commit (a4408e9) shipped without the trailer; an explicit follow-up trailer keeps the PR self-describing without rewriting published history. BREAKING CHANGE: Constructor `dim` default changes from `-1` (last axis) to `None` (full reduction) for SumFwdOp, MeanFwdOp, AmaxFwdOp, AminFwdOp, VarFwdOp, StdFwdOp, VarMeanFwdOp, AllFwdOp, AnyFwdOp, and CountNonzeroFwdOp -- the ten ops whose manifest entry declares `default: null`. ProdFwdOp preserves `dim=-1` (manifest declares `default: -1`). Callers depending on the old default MUST pass `dim=-1` explicitly. Per design, no deprecation cycle is introduced: this PR's mandate is strict alignment with the manifest spec, and the spec is authoritative under the project's design-first / spec-driven model.
lcy-seso
force-pushed
the
fix/reduction/issue-1434
branch
from
903a691 to
7a4e4dd
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 10 out of 10 changed files in this pull request and generated 2 comments.
Comments suppressed due to low confidence (1)
tileops/ops/reduction/_multidim.py:32
normalize_dim()is documented/used as acceptingdim=()(see docstring and tests), and the implementation supports tuples vialist(dim), but the type annotation is stillUnion[int, list[int], None]. Updating the annotation to includetuple[int, ...](or aSequence[int]) would keep the public helper signature consistent with actual supported inputs and the manifest’stuple[int, ...]form.
def normalize_dim(
dim: Union[int, list[int], None],
ndim: int,
*,
empty_dim_policy: EmptyDimPolicy = "reject",
) -> list[int]:
Ibuki-wind
left a comment
Ibuki-wind
left a comment
There was a problem hiding this comment.
Overall
One remaining reduction API-contract mismatch needs to be closed before this can safely unblock the status-flip PRs.
lcy-seso
changed the title
Co-Authored-By: Ibuki 🍃 — a wind born from GPTs <[email protected]>
Co-Authored-By: Ibuki 🍃 — a wind born from GPTs <[email protected]>
Ibuki-wind
left a comment
Ibuki-wind
left a comment
There was a problem hiding this comment.
Overall
One constructor contract change still breaks the current reduction test surface.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 11 out of 11 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (1)
tileops/ops/reduction/_multidim.py:33
normalize_dim()is now documented/used withdim=()(tuple) and returns[]for the newempty_dim_policy="noop", but its type annotation still only allowsUnion[int, list[int], None]. Please widen the annotation to includetuple[int, ...](or aSequence[int]), and consider explicitly rejectingbool(sinceisinstance(True, int)is true) sodim=True/dim=[True]can’t be accepted as a valid dimension spec.
def normalize_dim(
dim: Union[int, list[int], None],
ndim: int,
*,
empty_dim_policy: EmptyDimPolicy = "reject",
) -> list[int]:
"""Normalize and validate a dim specification.
Co-Authored-By: Ibuki 🍃 — a wind born from GPTs <[email protected]>
Ibuki-wind
left a comment
Ibuki-wind
left a comment
There was a problem hiding this comment.
Overall
Benchmarks still rely on the old implicit last-dimension default, so this PR can report invalid benchmark comparisons until those callsites are migrated.
Co-Authored-By: Ibuki 🍃 — a wind born from GPTs <[email protected]>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 13 out of 13 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (1)
tileops/ops/reduction/_multidim.py:32
normalize_dim()is documented/used withdim=()(tuple) and upstream call sites passTuple[int, ...], but its type annotation still only allowslist[int]and it will also acceptbool(sinceboolis anintsubclass) and silently treatTrueas axis 1. Please widen thedimtype to includetuple[int, ...](orSequence[int]) and add an explicitboolrejection (e.g., guard before theisinstance(dim, int)branch) to match the op-layerdimvalidation and avoid surprising behavior in other users like_softmax_base.
EmptyDimPolicy = Literal["reject", "full", "noop"]
def normalize_dim(
dim: Union[int, list[int], None],
ndim: int,
*,
empty_dim_policy: EmptyDimPolicy = "reject",
) -> list[int]:
Closes #1434
Constructor
dimdefault flips from-1toNonefor the ten ops whose manifest declaresdefault: null:SumFwdOp,MeanFwdOp,AmaxFwdOp,AminFwdOp,VarFwdOp,StdFwdOp,VarMeanFwdOp,AllFwdOp,AnyFwdOp,CountNonzeroFwdOp.ProdFwdOpkeepsdim=-1and narrows the type toint(its manifest declaresdim: int, default: -1).ProdFwdOp(dim=None | list | tuple)raisesTypeError.Callers depending on the old default must pass
dim=-1explicitly. No deprecation cycle — the flip is the spec alignment.Summary
dim=-1 → Nonein_ReduceOpBase,_WelfordReduceOp,AllFwdOp,AnyFwdOp,CountNonzeroFwdOp.EmptyDimPolicyextended toLiteral["reject", "full", "noop"].AllFwdOp/AnyFwdOpuse"noop", sodim=[]anddim=()short-circuit tooutput.shape == input.shape,dtype == bool,values == x.bool(). The short-circuit runs CUDA / dtype / ndim validation and binds_last_roofline_mn=(numel, 1)so the roofline still accounts for the read+write data movement.ProdFwdOp.__init__narrowed todim: int = -1with a_validate_dimoverride rejecting non-int._validate_dimrejectsbool(scalar and inside list/tuple) — a boolean is never a valid axis.dimannotation on the other constructors widened toUnion[int, List[int], Tuple[int, ...], None]to match the existingnormalize_dimruntime path.tileops/manifest/ortileops/kernels/.Test plan
pytest tests/ops/test_reduction*.py tests/ops/test_logical_reduce.py tests/ops/test_reduce.py tests/ops/test_reduce_dim_none.py tests/ops/test_reduce_multidim.py tests/ops/test_welford_non_aligned.py— green on CUDA.python scripts/validate_manifest.pyexits 0;tests/test_validate_manifest.py— 218 passed.git diff upstream/main -- tileops/{manifest,kernels}/— empty.pytest tests/test_validate_manifest.pyreturns 218 passed.New tests in
tests/ops/test_reduction_defaults.py(25 nodes) cover: per-op default-reduction shape,dim=[]/dim=()noop for All/Any,EmptyDimPolicy/normalize_dim/ class-policy bindings, noop input validation + roofline binding,ProdFwdOpdim: intenforcement, baseboolrejection.