Skip to content

Comments

ENH: add axis tuple support to torch.expand_dims#385

Merged
ev-br merged 1 commit intodata-apis:mainfrom
ev-br:expand_dims_tuples
Feb 23, 2026
Merged

ENH: add axis tuple support to torch.expand_dims#385
ev-br merged 1 commit intodata-apis:mainfrom
ev-br:expand_dims_tuples

Conversation

@ev-br
Copy link
Member

@ev-br ev-br commented Feb 5, 2026

  1. update for default axis (the argument is required)
  2. copycat what numpy does for the tuple axes
  3. TODO check the defaults for other wrappers (numpy,cupy,dask) all seem fine

Copilot AI review requested due to automatic review settings February 5, 2026 15:10
@ev-br ev-br marked this pull request as draft February 5, 2026 15:10
Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull request overview

This PR adds support for tuple values in the axis parameter of torch.expand_dims, allowing multiple dimensions to be expanded at once. This aligns with the Array API 2024.12 standard.

Changes:

  • Extended expand_dims function signature to accept axis: int | tuple[int, ...] instead of just axis: int
  • Added implementation logic to handle tuple axes by computing the expected output shape and using torch.reshape
  • Added validation for duplicate axes and out-of-range axis values

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

@ev-br ev-br force-pushed the expand_dims_tuples branch from 1f34958 to 35b631f Compare February 5, 2026 19:38
@ev-br ev-br added this to the 1.14 milestone Feb 5, 2026
@ev-br ev-br changed the title WIP: add axis tuple support to torch.expand_dims ENH: add axis tuple support to torch.expand_dims Feb 23, 2026
@ev-br ev-br marked this pull request as ready for review February 23, 2026 13:19
@ev-br ev-br merged commit 43de4af into data-apis:main Feb 23, 2026
28 checks passed
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