Skip to content

[torchlib] Implement aten::cdist and aten::_cdist_forward - #3171

Open
Qiming Bao (14H034160212) wants to merge 1 commit into
microsoft:mainfrom
14H034160212:torchlib-cdist
Open

Qiming Bao (14H034160212) wants to merge 1 commit into
microsoft:mainfrom
14H034160212:torchlib-cdist

Conversation

@14H034160212

Copy link
Copy Markdown

aten_cdist is still a NotImplementedError stub, so torch.onnx.export(..., dynamo=True) fails on any model that calls torch.cdist:

DispatchError: No ONNX function found for <OpOverload(op='aten._cdist_forward', overload='default')>

(#1647 was closed, but cdist was not implemented.)

This PR implements it for aten::cdist and aten::_cdist_forward, following PyTorch's semantics:

  • p == 2: uses the matmul form sqrt(max(|x1|² + |x2|²ᵀ − 2·x1·x2ᵀ, 0)), as PyTorch does, for compute_mode=1, and for the default mode when P or R is greater than 25 or not known statically. This avoids materializing the [*, P, R, M] difference tensor for large inputs. Otherwise, and for compute_mode=2, it uses the exact difference form.
  • Other p: reduces |x1[..., :, None, :] − x2[..., None, :, :]| over the last axis with ReduceL1 (p=1), ReduceL2, ReduceMax/ReduceMin (±inf), a non-zero count (p=0) or Pow/ReduceSum.
  • Batch dimensions broadcast as in PyTorch.

Tests:

  • Added cdist to ops_test_data.py. All 120 non-empty OpInfo subtests pass (p ∈ {1, 1.5, 2, 2.5, 3, inf}, both compute modes, broadcasting batch dims). Size-0 samples are skipped with a fixme, like flip and pixel_unshuffle: ORT returns the wrong shape when reducing a size-0 input.
  • End-to-end through torch.onnx.export(dynamo=True) + ONNX Runtime: 146/146 cases match eager (p ∈ {0, 1, 1.5, 2, 3, inf} × three compute modes × small/large/broadcast shapes × float32/float64, plus dynamic P and R).
  • ruff check / ruff format (0.15.8) clean.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Development

Successfully merging this pull request may close these issues.

1 participant