Skip to content

fix: handle initial and dtype in reductions - #74

Open
Robertboy18 wants to merge 1 commit into
aws-neuron:mainfrom
Robertboy18:fix/reduction-initial-dtype
Open

fix: handle initial and dtype in reductions#74
Robertboy18 wants to merge 1 commit into
aws-neuron:mainfrom
Robertboy18:fix/reduction-initial-dtype

Conversation

@Robertboy18

Copy link
Copy Markdown
Contributor

Summary

  • cast reduction inputs when dtype is specified
  • use the caller-provided initial value instead of always using the identity
  • validate tensor and array initializers are scalar
  • add coverage for sum, prod, max, and min initial values and explicit reduction dtypes

Problem

The HLO reduction helper accepted initial and dtype but ignored both. As a result, np.sum(x, initial=10) still initialized the reduction with zero, and np.sum(x, dtype=np.float16) still returned float32. The resulting HLO compiled successfully but had different semantics from NumPy.

Testing

  • pytest -q tests/unit/test_tensor_api.py -k "reduction or prod" (30 passed)
  • compared emitted HLO values, shapes, and dtypes with NumPy across 48 cases
  • compiled the regression kernels for trn2 with NeuronX Compiler 2.26

I did not have Trainium hardware available, so validation covered HLO semantics and compilation rather than device execution.

@Robertboy18
Robertboy18 requested review from a team and Copilot July 26, 2026 18:05

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@Robertboy18
Robertboy18 force-pushed the fix/reduction-initial-dtype branch from fd3d5ed to fdbb3cc Compare July 31, 2026 22:13
@ymwangg ymwangg added the blocked label Aug 4, 2026
@ymwangg

ymwangg commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

neuronx-cc currently doesn't support init value in reduce op, opened a ticket to compiler team NCC-9303.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants