Skip to content

Add torch.scatter#28

Merged
samuela merged 2 commits into
samuela:mainfrom
mavenlin:scatter
Sep 2, 2025
Merged

Add torch.scatter#28
samuela merged 2 commits into
samuela:mainfrom
mavenlin:scatter

Conversation

@mavenlin

@mavenlin mavenlin commented Aug 25, 2025

Copy link
Copy Markdown
Contributor
  • Aligned with the forward behavior of torch.scatter.
  • Backward only works for index.shape == src.shape, due to torch's own problem.
  • Add support to pass np.array directly for testing, because torch requests int64 for the index.
  • Also add support for Torchish to take np.ndarray as per our discussion Torchish.__init__ take np.ndarray? #23.

@samuela samuela left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

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

thanks @mavenlin ! overall lgtm

Comment thread torch2jax/__init__.py Outdated
Comment thread tests/test_core.py
@samuela
samuela merged commit 505769d into samuela:main Sep 2, 2025
1 check passed
@samuela

samuela commented Sep 2, 2025

Copy link
Copy Markdown
Owner

thanks @mavenlin !

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.

2 participants