Skip to content

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

Draft
ev-br wants to merge 1 commit intodata-apis:mainfrom
ev-br:expand_dims_tuples
Draft

WIP: add axis tuple support to torch.expand_dims#385
ev-br wants to merge 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) seems fine

Draft until 2025.12 is a thing and includes data-apis/array-api#988

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
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