apple / apple/coremltools

Support negative indices for pytorch.swapaxes

Open
#1,783 3 comments 0 reactions 0 assignees View on GitHub
feature request PyTorch (traced)
Dominant language
Python
Stars
5.4k
Forks
850
Avg merge
4d 5h
Merged PRs (30d)
10

Description

## 🌱 Describe your Feature Request
Please add support for negative indices in the Pytorch op `swapaxes`.

## How can this feature be used?
Swapaxes is a basic operation. Coremltools already supports it, but only with positive indices.

## Describe alternatives you've considered
- Use positive indices
- Use `transpose()`

## Additional context
Test code:

```python
import torch
import coremltools as ct

class MyModel(torch.nn.Module):
def __init__(self):
super().__init__()

def forward(self, x):
return x.swapaxes(-2, -1)

model = MyModel().eval()

with torch.inference_mode():
example_input = torch.rand(10, 20)
traced_model = torch.jit.trace(model, example_input)

coreml_model = ct.convert(
traced_model,
convert_to="neuralnetwork",
inputs=[ct.TensorType(shape=example_input.shape, name="x")],
debug=True,
)
```

Contributor guide

Open the contributing guide

Research direction

Locate the PyTorch swapaxes conversion entry point and the existing handling for axis indices; use the supplied traced-model example as the starting reproduction. Done means conversion of swapaxes(-2, -1) succeeds for the shown input, with a regression test covering the behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, tooling
Issue type
Feature
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.