Skip to content

feat: align tensordot with PyTorch Python API#812

Draft
voltjia wants to merge 1 commit into
masterfrom
feat/align-tensordot-python-api
Draft

feat: align tensordot with PyTorch Python API#812
voltjia wants to merge 1 commit into
masterfrom
feat/align-tensordot-python-api

Conversation

@voltjia

@voltjia voltjia commented Jul 24, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Add exactly two PyTorch-compatible tensordot(a, b, dims, out) forms: integer dims and a pair of dimension lists.
  • Keep the previous split dims_self, dims_other interface as a [[deprecated]] compatibility overload.

Motivation

The previous interface only exposed the lower-level pair of dimension lists. PyTorch exposes one dims argument that accepts either an integer or a pair of lists. The earlier revision also added a no-dims convenience overload; this revision removes it to keep the overload set minimal.

No standalone issue is associated with this alignment follow-up.

Type of Change

  • feat - new feature / new operator / new platform
  • fix - bug fix
  • perf - performance improvement (no behavioral change)
  • refactor - code restructuring without behavior change
  • test - adding or fixing tests only
  • docs - documentation only
  • build / ci - build system or CI configuration
  • chore - tooling, formatting, or other non-code changes
  • Breaking change

Platforms Affected

  • CPU (WITH_CPU)
  • NVIDIA (WITH_NVIDIA)
  • Iluvatar (WITH_ILUVATAR)
  • MetaX (WITH_METAX)
  • Cambricon (WITH_CAMBRICON)
  • Moore (WITH_MOORE)
  • Ascend (WITH_ASCEND)
  • PyTorch C++ bindings (WITH_TORCH)
  • Build system / CMake / CI
  • Python bindings / user-facing API

Smoke Test Result

Host: ssh nvidia
Image: accelerator-dev/nvidia:latest
PyTorch: 2.10.0a0+b4e4ee8

cmake ... -DWITH_CPU=ON -DWITH_NVIDIA=ON -DWITH_TORCH=ON \
  -DINFINI_OPS_OPS=tensordot -DINFINI_OPS_TORCH_OPS=tensordot
cmake --build build-nvidia --target ops -j 16
Build passed; codegen emitted 1 ATen implementation overload.

python -m pytest tests/test_generate_wrappers.py -q
13 passed

python -m pytest tests/test_tensordot.py --devices cpu nvidia -q
36 passed

Test Results on Supported Platforms

Platform Affected Build / Smoke Result Full Result / Notes
CPU Yes targeted build passed Integer, nested-list, and deprecated forms passed
NVIDIA Yes targeted build passed Integer, nested-list, and deprecated forms passed
Iluvatar No N/A - not affected N/A - not affected
MetaX No N/A - not affected N/A - not affected
Cambricon No N/A - not affected N/A - not affected
Moore No N/A - not affected N/A - not affected
Ascend No N/A - not affected N/A - not affected

Benchmark / Performance Impact

N/A. Integer dimensions are expanded during operator construction; the existing ATen execution path is unchanged.

Notes for Reviewers

API alignment

Interface Signature Alignment basis
InfiniOps integer form tensordot(a, b, dims, out) where dims is int64_t Matches one PyTorch public dims form; explicit out remains last per CONTRIBUTING.md.
InfiniOps list form tensordot(a, b, dims, out) where dims is two dimension lists Matches the other PyTorch public dims form.
PyTorch public API torch.tensordot(a, b, dims=2, out=None) Pinned PyTorch source
Deprecated compatibility tensordot(input, other, dims_self, dims_other, out) Preserves the pre-alignment split-list interface for migration.
  • C++ callers pass dims=2 explicitly when they want the PyTorch default.
  • No general wrapper-generator change is added by this PR.
  • Ruff 0.15.22 and clang-format 21.1.8 checks passed.

@voltjia
voltjia changed the base branch from master to fix/pin-ruff-version July 24, 2026 07:54
@voltjia
voltjia force-pushed the feat/align-tensordot-python-api branch from 5056ac1 to 94bda51 Compare July 24, 2026 08:02
Base automatically changed from fix/pin-ruff-version to master July 24, 2026 08:13
@voltjia
voltjia force-pushed the feat/align-tensordot-python-api branch from 94bda51 to 9f047f1 Compare July 24, 2026 08:27
@voltjia
voltjia force-pushed the feat/align-tensordot-python-api branch from 9f047f1 to 2360ed0 Compare July 24, 2026 11:07
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