Skip to content

Fix SliceUpdate JVPs with one traced input - #4200

Open
PhysicistJohn wants to merge 2 commits into
ml-explore:mainfrom
PhysicistJohn:fix-sliceupdate-single-input-jvp
Open

Fix SliceUpdate JVPs with one traced input#4200
PhysicistJohn wants to merge 2 commits into
ml-explore:mainfrom
PhysicistJohn:fix-sliceupdate-single-input-jvp

Conversation

@PhysicistJohn

Copy link
Copy Markdown
Contributor

Proposed changes

SliceUpdate::jvp currently rejects transforms when only the source or update input is traced. Initialize the absent tangent with zeros so assignment and sum slice updates support source-only and update-only JVPs while preserving the existing two-input behavior. Product, minimum, and maximum JVPs remain unsupported.

The focused regression passes 4/4 assertions. The full native suite passes 249/249 tests with 3,346 assertions.

Checklist

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

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