diff --git a/examples/kernels/flake.nix b/examples/kernels/flake.nix index 28498e74..8a136468 100644 --- a/examples/kernels/flake.nix +++ b/examples/kernels/flake.nix @@ -102,6 +102,12 @@ path = ./relu-torch-stable-abi; drv = sys: out: out.packages.${sys}.redistributable.${"torch-stable-abi211-${cudaVersion}-${sys}"}; } + { + # Exercise the shim with Torch 2.11 headers, which lack index_select. + name = "index-select-torch-stable-abi-kernel"; + path = ./index-select-torch-stable-abi; + drv = sys: out: out.packages.${sys}.redistributable.${"torch-stable-abi210-cu128-${sys}"}; + } # Check that we can also build the non-stable ABI version preceding the stable ABI version. { name = "relu-torch-stable-abi-kernel"; diff --git a/examples/kernels/index-select-torch-stable-abi/build.toml b/examples/kernels/index-select-torch-stable-abi/build.toml new file mode 100644 index 00000000..2c29baeb --- /dev/null +++ b/examples/kernels/index-select-torch-stable-abi/build.toml @@ -0,0 +1,23 @@ +[general] +name = "index-select-torch-stable-abi" +version = 1 +edition = 6 +license = "Apache-2.0" +backends = ["cuda"] + +[general.hub] +repo-id = "kernels-test/index-select-torch-stable-abi" + +[torch] +src = [ + "torch-ext/torch_binding.cpp", + "torch-ext/torch_binding.h", +] + +[torch.stable-abi] +cuda = "2.10" + +[kernel.index_select] +language = "cpp-cuda" +depends = ["torch", "torch-stable-abi-shim"] +src = ["index_select_cuda/index_select.cu"] diff --git a/examples/kernels/index-select-torch-stable-abi/flake.nix b/examples/kernels/index-select-torch-stable-abi/flake.nix new file mode 100644 index 00000000..db0fcdfa --- /dev/null +++ b/examples/kernels/index-select-torch-stable-abi/flake.nix @@ -0,0 +1,17 @@ +{ + description = "Flake for the Torch stable ABI index_select example"; + + inputs = { + kernel-builder.url = "path:../../.."; + }; + + outputs = + { + self, + kernel-builder, + }: + kernel-builder.lib.genKernelFlakeOutputs { + inherit self; + path = ./.; + }; +} diff --git a/examples/kernels/index-select-torch-stable-abi/index_select_cuda/index_select.cu b/examples/kernels/index-select-torch-stable-abi/index_select_cuda/index_select.cu new file mode 100644 index 00000000..23328ea9 --- /dev/null +++ b/examples/kernels/index-select-torch-stable-abi/index_select_cuda/index_select.cu @@ -0,0 +1,9 @@ +#include +#include + +torch::stable::Tensor index_select_impl( + const torch::stable::Tensor& input, + int64_t dim, + const torch::stable::Tensor& index) { + return torch_stable_abi_shim::index_select(input, dim, index); +} diff --git a/examples/kernels/index-select-torch-stable-abi/tests/test_index_select.py b/examples/kernels/index-select-torch-stable-abi/tests/test_index_select.py new file mode 100644 index 00000000..1bdf2e83 --- /dev/null +++ b/examples/kernels/index-select-torch-stable-abi/tests/test_index_select.py @@ -0,0 +1,25 @@ +import kernels +import pytest +import torch + +index_select_torch_stable_abi = kernels.get_kernel( + "kernels-test/index-select-torch-stable-abi", version=1 +) + + +@pytest.mark.kernels_ci +@pytest.mark.parametrize("dim", [0, 1, -1]) +@pytest.mark.parametrize("contiguous", [True, False]) +@pytest.mark.parametrize("indices", [[2, 0, 2], []]) +def test_index_select(device, dim, contiguous, indices): + if torch.device(device).type != "cuda": + pytest.skip("This example requires CUDA") + + input = torch.arange(24, device=device, dtype=torch.float32).reshape(4, 6) + if not contiguous: + input = input.t() + index = torch.tensor(indices, device=device, dtype=torch.int64) + + actual = index_select_torch_stable_abi.index_select(input, dim, index) + expected = torch.index_select(input, dim, index) + torch.testing.assert_close(actual, expected) diff --git a/examples/kernels/index-select-torch-stable-abi/torch-ext/index_select_torch_stable_abi/__init__.py b/examples/kernels/index-select-torch-stable-abi/torch-ext/index_select_torch_stable_abi/__init__.py new file mode 100644 index 00000000..384cb2fb --- /dev/null +++ b/examples/kernels/index-select-torch-stable-abi/torch-ext/index_select_torch_stable_abi/__init__.py @@ -0,0 +1,10 @@ +import torch + +from ._ops import ops + + +def index_select(input: torch.Tensor, dim: int, index: torch.Tensor) -> torch.Tensor: + return ops.index_select(input, dim, index) + + +__all__ = ["index_select"] diff --git a/examples/kernels/index-select-torch-stable-abi/torch-ext/torch_binding.cpp b/examples/kernels/index-select-torch-stable-abi/torch-ext/torch_binding.cpp new file mode 100644 index 00000000..864f78a4 --- /dev/null +++ b/examples/kernels/index-select-torch-stable-abi/torch-ext/torch_binding.cpp @@ -0,0 +1,14 @@ +#include + +#include "registration.h" +#include "torch_binding.h" + +STABLE_TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { + ops.def("index_select(Tensor input, int dim, Tensor index) -> Tensor"); +} + +STABLE_TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, ops) { + ops.impl("index_select", TORCH_BOX(&index_select_impl)); +} + +REGISTER_EXTENSION(TORCH_EXTENSION_NAME) diff --git a/examples/kernels/index-select-torch-stable-abi/torch-ext/torch_binding.h b/examples/kernels/index-select-torch-stable-abi/torch-ext/torch_binding.h new file mode 100644 index 00000000..d1bb3941 --- /dev/null +++ b/examples/kernels/index-select-torch-stable-abi/torch-ext/torch_binding.h @@ -0,0 +1,9 @@ +#pragma once + +#include +#include + +torch::stable::Tensor index_select_impl( + const torch::stable::Tensor& input, + int64_t dim, + const torch::stable::Tensor& index); diff --git a/kernels-common/src/config/deps.rs b/kernels-common/src/config/deps.rs index aea871df..dfd3fa6b 100644 --- a/kernels-common/src/config/deps.rs +++ b/kernels-common/src/config/deps.rs @@ -34,6 +34,8 @@ pub enum Dependency { #[serde(rename = "metal-cpp")] MetalCpp, Torch, + #[serde(rename = "torch-stable-abi-shim")] + TorchStableAbiShim, } #[derive(Debug, Deserialize, Serialize)] @@ -102,3 +104,16 @@ pub enum DependencyError { #[error("Unknown dependency: `{dependency:?}`")] GeneralDependency { dependency: String }, } + +#[cfg(test)] +mod tests { + use super::Dependency; + + #[test] + fn torch_stable_abi_shim_round_trips() { + let json = r#""torch-stable-abi-shim""#; + let dependency: Dependency = serde_json::from_str(json).unwrap(); + assert_eq!(dependency, Dependency::TorchStableAbiShim); + assert_eq!(serde_json::to_string(&dependency).unwrap(), json); + } +} diff --git a/nix-builder/lib/cpp-deps.nix b/nix-builder/lib/cpp-deps.nix index f5c38ec0..416fbc23 100644 --- a/nix-builder/lib/cpp-deps.nix +++ b/nix-builder/lib/cpp-deps.nix @@ -34,6 +34,9 @@ let "torch" = [ torch ]; + "torch-stable-abi-shim" = [ + pkgs.torch-stable-abi-shim + ]; "sycl_tla" = [ (torch.xpuPackages.sycl-tla.override { inherit stdenv; }) ]; diff --git a/nix-builder/overlay.nix b/nix-builder/overlay.nix index b41df3f2..905b5d8f 100644 --- a/nix-builder/overlay.nix +++ b/nix-builder/overlay.nix @@ -59,6 +59,8 @@ final: prev: torch-ops-check = final.callPackage ./pkgs/torch-ops-check { }; + torch-stable-abi-shim = final.callPackage ./pkgs/torch-stable-abi-shim { }; + ruff = prev.ruff.overrideAttrs ( finalAttrs: prevAttrs: { version = "0.15.10"; diff --git a/nix-builder/pkgs/torch-stable-abi-shim/default.nix b/nix-builder/pkgs/torch-stable-abi-shim/default.nix new file mode 100644 index 00000000..c85fe00c --- /dev/null +++ b/nix-builder/pkgs/torch-stable-abi-shim/default.nix @@ -0,0 +1,27 @@ +{ + lib, + stdenvNoCC, +}: + +stdenvNoCC.mkDerivation { + name = "torch-stable-abi-shim"; + + src = ../../../torch-stable-abi-shim; + + dontConfigure = true; + dontBuild = true; + + installPhase = '' + runHook preInstall + + mkdir -p "$out/include" + cp -r include/. "$out/include/" + + runHook postInstall + ''; + + meta = { + description = "Header-only shim for the PyTorch stable ABI"; + platforms = lib.platforms.all; + }; +} diff --git a/torch-stable-abi-shim/LICENSE b/torch-stable-abi-shim/LICENSE new file mode 100644 index 00000000..c23172f7 --- /dev/null +++ b/torch-stable-abi-shim/LICENSE @@ -0,0 +1,84 @@ +From PyTorch: + +Copyright (c) 2016- Facebook, Inc (Adam Paszke) +Copyright (c) 2014- Facebook, Inc (Soumith Chintala) +Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert) +Copyright (c) 2012-2014 Deepmind Technologies (Koray Kavukcuoglu) +Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu) +Copyright (c) 2011-2013 NYU (Clement Farabet) +Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston) +Copyright (c) 2006 Idiap Research Institute (Samy Bengio) +Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz) + +From Caffe2: + +Copyright (c) 2016-present, Facebook Inc. All rights reserved. + +All contributions by Facebook: +Copyright (c) 2016 Facebook Inc. + +All contributions by Google: +Copyright (c) 2015 Google Inc. +All rights reserved. + +All contributions by Yangqing Jia: +Copyright (c) 2015 Yangqing Jia +All rights reserved. + +All contributions by Kakao Brain: +Copyright 2019-2020 Kakao Brain + +All contributions by Cruise LLC: +Copyright (c) 2022 Cruise LLC. +All rights reserved. + +All contributions by Tri Dao: +Copyright (c) 2024 Tri Dao. +All rights reserved. + +All contributions by Arm: +Copyright (c) 2021, 2023-2025 Arm Limited and/or its affiliates + +All contributions from Caffe: +Copyright(c) 2013, 2014, 2015, the respective contributors +All rights reserved. + +All other contributions: +Copyright(c) 2015, 2016 the respective contributors +All rights reserved. + +Caffe2 uses a copyright model similar to Caffe: each contributor holds +copyright over their contributions to Caffe2. The project versioning records +all such contribution and copyright details. If a contributor wants to further +mark their specific copyright on a particular contribution, they should +indicate their copyright solely in the commit message of the change when it is +committed. + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + +3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America + and IDIAP Research Institute nor the names of its contributors may be + used to endorse or promote products derived from this software without + specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE +LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +POSSIBILITY OF SUCH DAMAGE. diff --git a/torch-stable-abi-shim/README.md b/torch-stable-abi-shim/README.md new file mode 100644 index 00000000..d088238d --- /dev/null +++ b/torch-stable-abi-shim/README.md @@ -0,0 +1,15 @@ +# Torch stable ABI shim + +A header-only shim that makes newer PyTorch stable C++ API wrappers available +when building against older PyTorch stable headers. Wrappers introduced in later +headers can use the stable dispatcher API already available in older versions. +APIs requiring a newer runtime are guarded by `TORCH_FEATURE_VERSION` and are +only exposed when the targeted stable ABI version supports them. + +The headers are vendored from PyTorch, with names adjusted to avoid collisions +with PyTorch. Tensor and device types remain in `torch::stable`. + +## Copyright and license + +See [LICENSE](LICENSE) for the unmodified upstream PyTorch license, including +all copyright notices, conditions, and the disclaimer. diff --git a/torch-stable-abi-shim/include/torch_stable_abi_shim/csrc/stable/macros.h b/torch-stable-abi-shim/include/torch_stable_abi_shim/csrc/stable/macros.h new file mode 100644 index 00000000..99acd15f --- /dev/null +++ b/torch-stable-abi-shim/include/torch_stable_abi_shim/csrc/stable/macros.h @@ -0,0 +1,172 @@ +#pragma once + +// Vendored from PyTorch commit d3c397fef74cebde85dde5ca5e6638a2cb2da9c5. +// Helpers and macros are namespaced to coexist with PyTorch headers. +#include +#include + +#include +#include +#include +#include +#include + +#if defined(_WIN32) +#include +#else +#include +#endif + +#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + +// Users of this macro are expected to include cuda_runtime.h +#define TORCH_STABLE_ABI_SHIM_CUDA_CHECK(EXPR) \ + do { \ + const cudaError_t __err = EXPR; \ + char* __error_msg = nullptr; \ + torch_c10_cuda_check_msg( \ + static_cast(__err), \ + __FILE__, \ + __func__, \ + static_cast(__LINE__), \ + true, \ + &__error_msg); \ + if (__error_msg != nullptr) { \ + std::string __msg(__error_msg); \ + torch_c10_cuda_free_error_msg(__error_msg); \ + /* @allow-raw-throw: the message already names the failing call */ \ + throw std::runtime_error(__msg); \ + } \ + } while (0) + +// Users of this macro are expected to include cuda_runtime.h +#define TORCH_STABLE_ABI_SHIM_CUDA_KERNEL_LAUNCH_CHECK() TORCH_STABLE_ABI_SHIM_CUDA_CHECK(cudaGetLastError()) + +#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + +namespace torch_stable_abi_shim::detail { +// Look up a symbol already loaded in the current process. Used to reach +// stable-ABI shims that were added after the extension's TORCH_TARGET_VERSION +// when the running libtorch is new enough to provide them. Returns nullptr if +// the symbol is absent or the platform has no in-process dynamic lookup. This +// is the single place that holds the per-platform lookup; the macros below are +// platform agnostic. +[[maybe_unused]] static void* lookup_stable_symbol(const char* name) { +#if defined(C10_MOBILE) + (void)name; + return nullptr; +#elif defined(_WIN32) + // Windows has no RTLD_DEFAULT equivalent, so resolve against the modules that + // export the stable shims. Today these live in torch_cpu.dll and + // torch_cuda.dll, so try each in turn. GetModuleHandleW does not take a + // reference on the returned handle. + const wchar_t* dlls[] = {L"torch_cpu.dll", L"torch_cuda.dll"}; + for (const wchar_t* dll : dlls) { + HMODULE handle = GetModuleHandleW(dll); + if (handle != nullptr) { + void* sym = reinterpret_cast(GetProcAddress(handle, name)); + if (sym != nullptr) { + return sym; + } + } + } + return nullptr; +#else + return dlsym(RTLD_DEFAULT, name); +#endif +} + +// Fallback for the const char* exception getters, returns a nullptr. A fallback +// MUST share the exact signature (return type and arguments) of the shim it +// backs: the looked-up pointer is called through decltype(&FALLBACK_FUNCTION). +[[maybe_unused]] static const char* torch_shim_bc_const_char_ptr() { + return nullptr; +} +} // namespace torch_stable_abi_shim::detail + +// IMPORTANT: this macro is intended for RARE, EXCEPTIONAL use only. It performs +// an at-runtime symbol lookup (dlsym / GetProcAddress), which is relatively +// expensive and only works on platforms that support dynamic symbol lookup. +// Reach for it only when an older-target extension genuinely needs a newer shim +// that would otherwise be unavailable (e.g. richer error messages); prefer +// normal version-guarded shim calls everywhere else. +// +// Calls SHIM_FUNCTION via the runtime lookup when the running libtorch is at +// least VERSION_SHIM_PRESENT, otherwise calls FALLBACK_FUNCTION. +// +// We structure this macro as an immediately-invoked lambda so it can 'return' a +// value from the macro call. The called pointer type is derived from +// FALLBACK_FUNCTION, which must therefore share SHIM_FUNCTION's exact signature +// (including arguments). +#define TORCH_STABLE_ABI_SHIM_DYNAMIC_VERSION_CALL( \ + VERSION_SHIM_PRESENT, SHIM_FUNCTION, FALLBACK_FUNCTION, ...) \ + ([&]() { \ + if (aoti_torch_abi_version() >= (VERSION_SHIM_PRESENT)) { \ + void* fn_ptr = \ + torch_stable_abi_shim::detail::lookup_stable_symbol(#SHIM_FUNCTION); \ + if (fn_ptr != nullptr) { \ + using FunctionType = decltype(&FALLBACK_FUNCTION); \ + return reinterpret_cast(fn_ptr)(__VA_ARGS__); \ + } \ + } \ + return FALLBACK_FUNCTION(__VA_ARGS__); \ + })() + +// Entry point for dynamic version calls for shims added in 2.13.0. Putting the +// version expectation in the macro lets us bypass the symbol lookup entirely +// when the target version already includes the shim and call it directly; +// otherwise we fall through to the runtime lookup. +#if TORCH_FEATURE_VERSION >= (((0ULL + 2) << 56) | ((0ULL + 13) << 48)) +// Target version already includes the shim, call it directly. +#define TORCH_STABLE_ABI_SHIM_DYNAMIC_VERSION_CALL_2_13_0( \ + SHIM_FUNCTION, FALLBACK_FUNCTION, ...) \ + ([&]() { return SHIM_FUNCTION(__VA_ARGS__); })() +#else +// Target version predates the shim, try a dynamic lookup. +#define TORCH_STABLE_ABI_SHIM_DYNAMIC_VERSION_CALL_2_13_0( \ + SHIM_FUNCTION, FALLBACK_FUNCTION, ...) \ + TORCH_STABLE_ABI_SHIM_DYNAMIC_VERSION_CALL( \ + (((0ULL + 2) << 56) | ((0ULL + 13) << 48)), SHIM_FUNCTION, FALLBACK_FUNCTION, __VA_ARGS__) +#endif // TORCH_FEATURE_VERSION >= (((0ULL + 2) << 56) | ((0ULL + 13) << 48)) + +namespace torch_stable_abi_shim::detail { +[[maybe_unused]] C10_NOINLINE static void throw_exception( + const char* call, + const char* file, + int64_t line) { + std::stringstream ss; + const auto& error_msg_without = TORCH_STABLE_ABI_SHIM_DYNAMIC_VERSION_CALL_2_13_0( + torch_exception_get_what_without_backtrace, torch_shim_bc_const_char_ptr); + if (error_msg_without) { + ss << error_msg_without; + ss << " (originally from " << call << " API call failed at " << file + << ", line " << line << ")"; + } else { + ss << call << " API call failed at " << file << ", line " << line; + } + + const auto& error_msg = TORCH_STABLE_ABI_SHIM_DYNAMIC_VERSION_CALL_2_13_0( + torch_exception_get_what, torch_shim_bc_const_char_ptr); + if (error_msg) { + std::time_t t = std::time(nullptr); + std::tm tm = *std::localtime(&t); + std::cerr << "[" << std::put_time(&tm, "%H:%M:%S") << " " << file << ":" + << line + << "] Exception across libtorch C API boundary: " << error_msg; + } + + // STD_TORCH_CHECK would stamp this helper's file and line over the + // caller's, which the message above already carries. + // @allow-raw-throw: keeps the caller's location in the message + throw std::runtime_error(ss.str()); +} +} // namespace torch_stable_abi_shim::detail + +// This macro is similar to the header-only macro TORCH_ERROR_CODE_CHECK, +// but this macro is NOT header-only! It depends on the stable ABI but +// provides more info in the exception, including the error message as +// retrieved through the c shims from the original error message. +#define TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(call) \ + if ((call) != TORCH_SUCCESS) { \ + torch_stable_abi_shim::detail::throw_exception(#call, __FILE__, __LINE__); \ + } diff --git a/torch-stable-abi-shim/include/torch_stable_abi_shim/csrc/stable/ops.h b/torch-stable-abi-shim/include/torch_stable_abi_shim/csrc/stable/ops.h new file mode 100644 index 00000000..8b1a2fa1 --- /dev/null +++ b/torch-stable-abi-shim/include/torch_stable_abi_shim/csrc/stable/ops.h @@ -0,0 +1,1628 @@ +#pragma once + +// Vendored from PyTorch commit d3c397fef74cebde85dde5ca5e6638a2cb2da9c5. +// APIs are gated by the targeted PyTorch stable ABI version. +// Encoded version checks also work with older headers lacking newer constants. +// Upstream header availability records when PyTorch shipped each wrapper; +// it is not a build-time requirement for this shim. + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace torch_stable_abi_shim { + +/// A function pointer type for data deleters used with from_blob. +/// The deleter is called with the data pointer when the tensor's storage +/// is deallocated. +using DeleterFnPtr = void (*)(void*); + +/// Stable version of the empty_like op. +/// +/// Creates a new uninitialized tensor with the same size, dtype, layout, and +/// device as the input tensor. This version does not support kwargs (device, +/// dtype, layout, memory_format) - kwargs support may be added in the future. +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @param self The input tensor whose properties will be used for the new +/// tensor. +/// @return A new uninitialized tensor with the same properties as self. +inline torch::stable::Tensor empty_like(const torch::stable::Tensor& self) { + const auto num_args = 6; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(std::nullopt), + torch::stable::detail::from(std::nullopt), + torch::stable::detail::from(std::nullopt), + torch::stable::detail::from(std::nullopt), + torch::stable::detail::from(std::nullopt)}; +#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::empty_like", "", stack.data(), TORCH_ABI_VERSION)); +#else + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_call_dispatcher("aten::empty_like", "", stack.data())); +#endif + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the fill_.Scalar op. +/// +/// Fills the input tensor with the specified scalar value in-place and returns +/// it. This has identical semantics to the existing fill_.Scalar op. +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @note The value parameter is typed as double +/// This is because Scalar.h is currently not header-only. +/// +/// @param self The tensor to fill. +/// @param value The scalar value to fill the tensor with. +/// @return The input tensor, now filled with the specified value. +inline torch::stable::Tensor fill_( + const torch::stable::Tensor& self, + double value) { + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_aten_fill__Scalar(self.get(), value)); + return self; +} + +/// Stable version of the narrow.default op. +/// +/// Returns a new tensor that is a narrowed version of the input tensor. The +/// dimension dim is narrowed from start to start + length. +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @note The start and length parameters +/// is not yet header-only. +/// +/// @param self The input tensor to narrow. +/// @param dim The dimension along which to narrow. +/// @param start The starting index for the narrowed dimension. +/// @param length The length of the narrowed dimension. +/// @return A new tensor that is a narrowed view of the input. +inline torch::stable::Tensor narrow( + torch::stable::Tensor& self, + int64_t dim, + int64_t start, + int64_t length) { + AtenTensorHandle ret0 = nullptr; + + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_aten_narrow(self.get(), dim, start, length, &ret0)); + return torch::stable::Tensor(ret0); +} + +#if TORCH_FEATURE_VERSION < TORCH_VERSION_2_10_0 +/// Stable version of the new_empty op (2.9 version). +/// +/// Creates a new uninitialized tensor with the specified size, inheriting +/// device and layout from the input tensor. This version only supports the +/// dtype kwarg. For the full kwargs version, use PyTorch 2.10+. +/// +/// Minimum compatible version: PyTorch 2.9. For full kwargs support, use +/// PyTorch 2.10+. +/// +/// @param self The input tensor whose device +/// @param size The desired size of the output tensor. +/// @param dtype Optional scalar type for the tensor elements. If not provided, +/// inherits from self. +/// @return A new uninitialized tensor with the specified properties. +inline torch::stable::Tensor new_empty( + const torch::stable::Tensor& self, + torch::headeronly::IntHeaderOnlyArrayRef size, + std::optional dtype = std::nullopt) { + int32_t device_type; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_get_device_type(self.get(), &device_type)); + + int32_t device_index; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_get_device_index(self.get(), &device_index)); + + int32_t target_dtype; + if (dtype.has_value()) { + target_dtype = torch::stable::detail::to( + torch::stable::detail::from(dtype.value())); + } else { + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_get_dtype(self.get(), &target_dtype)); + } + + int32_t layout; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(aoti_torch_get_layout(self.get(), &layout)); + + AtenTensorHandle ret0; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(aoti_torch_aten_new_empty( + self.get(), + size.data(), + static_cast(size.size()), + &target_dtype, + &layout, + &device_type, + device_index, + nullptr, // pin_memory (nullptr for default) + &ret0)); + + return torch::stable::Tensor(ret0); +} + +/// Stable version of the new_zeros op (2.9 version). +/// +/// Creates a new tensor filled with zeros with the specified size, inheriting +/// device and layout from the input tensor. This version only supports the +/// dtype kwarg. For the full kwargs version, use PyTorch 2.10+. +/// +/// Minimum compatible version: PyTorch 2.9. For full kwargs support, use +/// PyTorch 2.10+. +/// +/// @param self The input tensor whose device and layout will be inherited. +/// @param size The desired size of the output tensor. +/// @param dtype Optional scalar type for the tensor elements. If not provided, +/// inherits from self. +/// @return A new zero-filled tensor with the specified properties. +inline torch::stable::Tensor new_zeros( + const torch::stable::Tensor& self, + torch::headeronly::IntHeaderOnlyArrayRef size, + std::optional dtype = std::nullopt) { + int32_t device_type; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_get_device_type(self.get(), &device_type)); + + int32_t device_index; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_get_device_index(self.get(), &device_index)); + + int32_t target_dtype; + if (dtype.has_value()) { + target_dtype = torch::stable::detail::to( + torch::stable::detail::from(dtype.value())); + } else { + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_get_dtype(self.get(), &target_dtype)); + } + + int32_t layout; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(aoti_torch_get_layout(self.get(), &layout)); + + AtenTensorHandle ath; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(aoti_torch_aten_new_zeros( + self.get(), + size.data(), + static_cast(size.size()), + &target_dtype, + &layout, + &device_type, + device_index, + nullptr, // pin_memory (nullptr for default) + &ath)); + + return torch::stable::Tensor(ath); +} +#endif // TORCH_FEATURE_VERSION < TORCH_VERSION_2_10_0 + +/// Stable version of the pad.default op. +/// +/// Pads the input tensor according to the specified padding sizes. The padding +/// is applied symmetrically to each dimension, with the padding sizes specified +/// in reverse order (last dimension first). +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @note The pad parameter is typed +/// not yet header-only. +/// +/// @param self The input tensor to pad. +/// @param pad The padding sizes for each dimension (in pairs, starting from +/// the last dimension). +/// @param mode The padding mode: "constant", "reflect", "replicate", or +/// "circular". Defaults to "constant". +/// @param value The fill value for constant padding. Defaults to 0.0. +/// @return A new padded tensor. +inline torch::stable::Tensor pad( + const torch::stable::Tensor& self, + torch::headeronly::IntHeaderOnlyArrayRef pad, + const std::string& mode = "constant", + double value = 0.0) { + AtenTensorHandle ret0 = nullptr; + + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(aoti_torch_aten_pad( + self.get(), pad.data(), pad.size(), mode.c_str(), &value, &ret0)); + return torch::stable::Tensor(ret0); +} + +/// Stable version of the amax.default op (single dimension). +/// +/// Computes the maximum value along the specified dimension. If keepdim is +/// true, the output tensor has the same number of dimensions as the input, +/// with the reduced dimension having size 1. Otherwise, the reduced dimension +/// is removed. +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @param self The input tensor. +/// @param dim The dimension along which to compute the maximum. +/// @param keepdim Whether to retain +/// @return A tensor containing the maximum values along the specified +/// dimension. +inline torch::stable::Tensor amax( + const torch::stable::Tensor& self, + int64_t dim, + bool keepdim = false) { + AtenTensorHandle ret = nullptr; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_aten_amax(self.get(), &dim, 1, keepdim, &ret)); + return torch::stable::Tensor(ret); +} + +/// Stable version of the amax.default op (multiple dimensions). +/// +/// Computes the maximum value reducing over all the specified dimensions. If +/// keepdim is true, the output tensor has the same number of dimensions as the +/// input, with the reduced dimensions having size 1. Otherwise, the reduced +/// dimensions are removed. +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @note The dims parameter is typed +/// is not yet header-only. +/// +/// @param self The input tensor. +/// @param dims The dimensions along which to compute the maximum. +/// @param keepdim Whether to retain the reduced dimensions. Defaults to false. +/// @return A tensor containing the maximum values. +inline torch::stable::Tensor amax( + const torch::stable::Tensor& self, + torch::headeronly::IntHeaderOnlyArrayRef dims, + bool keepdim = false) { + AtenTensorHandle ret = nullptr; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(aoti_torch_aten_amax( + self.get(), + dims.data(), + static_cast(dims.size()), + keepdim, + &ret)); + return torch::stable::Tensor(ret); +} + +/// Stable version of the transpose.int op. +/// +/// Returns a tensor that is a transposed version of the input, with dimensions +/// dim0 and dim1 swapped. The returned tensor shares storage with the input. +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @param self The input tensor. +/// @param dim0 The first dimension to transpose. +/// @param dim1 The second dimension to transpose. +/// @return A transposed view of the input tensor. +inline torch::stable::Tensor transpose( + const torch::stable::Tensor& self, + int64_t dim0, + int64_t dim1) { + const auto num_args = 3; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(dim0), + torch::stable::detail::from(dim1)}; +#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::transpose", "int", stack.data(), TORCH_ABI_VERSION)); +#else + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_call_dispatcher("aten::transpose", "int", stack.data())); +#endif + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the zero_ op. +/// +/// Fills the input tensor with zeros in-place and returns it. Unlike the +/// tensor method version (t.zero_()), this is called as a function: zero_(t). +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @param self The tensor to fill with zeros. +/// @return The input tensor, now filled with zeros. +inline torch::stable::Tensor zero_(torch::stable::Tensor& self) { + const auto num_args = 1; + std::array stack{torch::stable::detail::from(self)}; +#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::zero_", "", stack.data(), TORCH_ABI_VERSION)); +#else + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_call_dispatcher("aten::zero_", "", stack.data())); +#endif + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the copy_ op. +/// +/// Copies the elements from the source tensor into the destination tensor +/// in-place and returns the destination tensor. The tensors must be +/// broadcastable. +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @param self The destination tensor (modified in-place). +/// @param src The source tensor to copy from. +/// @param non_blocking If true, the copy may occur asynchronously with respect +/// to the host. Defaults to false. +/// @return The destination tensor with copied values. +inline torch::stable::Tensor copy_( + torch::stable::Tensor& self, + const torch::stable::Tensor& src, + std::optional non_blocking = std::nullopt) { + const auto num_args = 3; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(src), + torch::stable::detail::from(non_blocking.value_or(false))}; +#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::copy_", "", stack.data(), TORCH_ABI_VERSION)); +#else + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_call_dispatcher("aten::copy_", "", stack.data())); +#endif + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the clone op. +/// +/// Returns a copy of the input tensor. The returned tensor has the same data +/// and type as the input, but is stored in a new memory location. +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @note Optional memory_format kwarg support +/// +/// @param self The input tensor to clone. +/// @return A new tensor with copied data. +inline torch::stable::Tensor clone(const torch::stable::Tensor& self) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(std::nullopt)}; +#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::clone", "", stack.data(), TORCH_ABI_VERSION)); +#else + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_call_dispatcher("aten::clone", "", stack.data())); +#endif + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the flatten.using_ints op. +/// +/// Flattens the input tensor by reshaping it into a one-dimensional tensor. +/// If start_dim or end_dim are specified, only dimensions starting from +/// start_dim to end_dim are flattened. +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @param self The input tensor to flatten. +/// @param start_dim The first dimension to flatten. Defaults to 0. +/// @param end_dim The last dimension to flatten. Defaults to -1 (last dim). +/// @return A flattened tensor. +inline torch::stable::Tensor flatten( + const torch::stable::Tensor& self, + int64_t start_dim = 0, + int64_t end_dim = -1) { + const auto num_args = 3; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(start_dim), + torch::stable::detail::from(end_dim)}; +#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::flatten", "using_ints", stack.data(), TORCH_ABI_VERSION)); +#else + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_call_dispatcher("aten::flatten", "using_ints", stack.data())); +#endif + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the unsqueeze op. +/// +/// Returns a new tensor with a dimension of size one inserted at the specified +/// position. The returned tensor shares the same underlying data with the input +/// tensor. +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @param self The input tensor. +/// @param dim The index at which to insert +/// values are supported. +/// @return A tensor with an additional dimension. +inline torch::stable::Tensor unsqueeze( + const torch::stable::Tensor& self, + int64_t dim) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(dim)}; +#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::unsqueeze", "", stack.data(), TORCH_ABI_VERSION)); +#else + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_call_dispatcher("aten::unsqueeze", "", stack.data())); +#endif + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the squeeze.dim op. +/// +/// Returns a tensor with the dimension of size one at the specified position +/// removed. The returned tensor shares the same underlying data with the input +/// tensor. +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @param self The input tensor. +/// @param dim The dimension to squeeze. +/// the tensor is returned unchanged. +/// @return A tensor with the specified dimension removed (if size was 1). +inline torch::stable::Tensor squeeze( + const torch::stable::Tensor& self, + int64_t dim) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(dim)}; +#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::squeeze", "dim", stack.data(), TORCH_ABI_VERSION)); +#else + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_call_dispatcher("aten::squeeze", "dim", stack.data())); +#endif + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the select.int op. +/// +/// Slices the input tensor along the specified dimension at the given index. +/// This function returns a view of the original tensor with the given dimension +/// removed. +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @note The index parameter is typed +/// header-only. +/// +/// @param self The input tensor. +/// @param dim The dimension to slice. +/// @param index The index to select along the dimension. +/// @return A tensor with one fewer dimension. +inline torch::stable::Tensor select( + const torch::stable::Tensor& self, + int64_t dim, + int64_t index) { + const auto num_args = 3; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(dim), + torch::stable::detail::from(index)}; +#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::select", "int", stack.data(), TORCH_ABI_VERSION)); +#else + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_call_dispatcher("aten::select", "int", stack.data())); +#endif + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the matmul op. +/// +/// Performs matrix multiplication between two tensors. The behavior depends on +/// the dimensionality of the tensors (see PyTorch documentation for details on +/// broadcasting rules for matmul). +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @param self The first input tensor. +/// @param other The second input tensor. +/// @return The result of matrix multiplication. +inline torch::stable::Tensor matmul( + const torch::stable::Tensor& self, + const torch::stable::Tensor& other) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(other)}; +#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::matmul", "", stack.data(), TORCH_ABI_VERSION)); +#else + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_call_dispatcher("aten::matmul", "", stack.data())); +#endif + return torch::stable::detail::to(stack[0]); +} + +#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + +/// Stable parallel_for utility. +/// +/// Provides a stable interface to at::parallel_for for parallel execution. +/// The function f will be called with (begin, end) ranges to process in +/// parallel. grain_size controls the minimum work size per thread for efficient +/// parallelization. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @tparam F The callable type +/// @param begin The start of the iteration range. +/// @param end The end of the iteration range (exclusive). +/// @param grain_size The minimum number of iterations per thread. +/// @param f The function to execute in parallel. +template +inline void parallel_for( + const int64_t begin, + const int64_t end, + const int64_t grain_size, + const F& f) { + auto callback = [](int64_t cb_begin, int64_t cb_end, void* ctx) { + const F* func = static_cast(ctx); + (*func)(cb_begin, cb_end); + }; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_parallel_for( + begin, + end, + grain_size, + callback, + const_cast(static_cast(&f)))); +} + +/// Gets the number of threads for the parallel backend. +/// +/// Provides a stable interface to at::get_num_threads. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @return The number of threads +inline uint32_t get_num_threads() { + uint32_t num_threads; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_get_num_threads(&num_threads)); + return num_threads; +} + +/// Stable version of the empty.memory_format op. +/// +/// Creates a new uninitialized tensor with the specified size and options. +/// This function supports full tensor creation options including device, +/// dtype, layout, and memory format. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @param size The desired size of the output tensor. +/// @param dtype Optional scalar type for the tensor elements. +/// @param layout Optional memory layout (e.g., strided, sparse). +/// @param device Optional device to place the tensor on. +/// @param pin_memory Optional flag to use pinned memory (for CUDA tensors). +/// @param memory_format Optional memory format for the tensor. +/// @return A new uninitialized tensor with the specified properties. +inline torch::stable::Tensor empty( + torch::headeronly::IntHeaderOnlyArrayRef size, + std::optional dtype = std::nullopt, + std::optional layout = std::nullopt, + std::optional device = std::nullopt, + std::optional pin_memory = std::nullopt, + std::optional memory_format = + std::nullopt) { + const auto num_args = 6; + std::array stack{ + torch::stable::detail::from(size), + torch::stable::detail::from(dtype), + torch::stable::detail::from(layout), + torch::stable::detail::from(device), + torch::stable::detail::from(pin_memory), + torch::stable::detail::from(memory_format)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::empty", "memory_format", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the reshape op. +/// +/// Returns a tensor with the same data and number of elements as the input, +/// but with the specified shape. When possible, the returned tensor will be +/// a view of the input. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @param self The input tensor. +/// @param shape The desired output shape. +/// @return A tensor with the specified shape. +inline torch::stable::Tensor reshape( + const torch::stable::Tensor& self, + torch::headeronly::IntHeaderOnlyArrayRef shape) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(shape)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::reshape", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the view op. +/// +/// Returns a new tensor with the same data as the input tensor but with a +/// different shape. The returned tensor shares the same data and must have +/// the same number of elements. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @param self The input tensor. +/// @param size The desired output shape. +/// @return A view tensor with the specified shape. +inline torch::stable::Tensor view( + const torch::stable::Tensor& self, + torch::headeronly::IntHeaderOnlyArrayRef size) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(size)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + torch_call_dispatcher("aten::view", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the view.dtype op. +/// +/// Returns a view of the input tensor with the same data reinterpreted as the +/// given dtype. This is a different overload than the shape-based view op. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.14. +/// +/// @param self The input tensor. +/// @param dtype The dtype to reinterpret the data as. +/// @return A view of the input tensor with the specified dtype. +inline torch::stable::Tensor view( + const torch::stable::Tensor& self, + torch::headeronly::ScalarType dtype) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(dtype)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::view", "dtype", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Creates a tensor from an existing data blob. +/// +/// Creates a tensor that uses the provided data pointer as its storage. +/// The tensor does not own the data, so the caller must ensure the data +/// remains valid for the lifetime of the tensor. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @param data Pointer to the data buffer. +/// @param sizes The size of each dimension of the tensor. +/// @param strides The stride for each dimension. +/// @param device The device where the data resides. +/// @param dtype The scalar type of the data. +/// @param storage_offset The offset into the data buffer. Defaults to 0. +/// @param layout The memory layout. Defaults to Strided. +/// @return A tensor backed by the provided data. +inline torch::stable::Tensor from_blob( + void* data, + torch::headeronly::IntHeaderOnlyArrayRef sizes, + torch::headeronly::IntHeaderOnlyArrayRef strides, + torch::stable::Device device, + torch::headeronly::ScalarType dtype, + int64_t storage_offset = 0, + torch::headeronly::Layout layout = torch::headeronly::Layout::Strided) { + auto shim_dtype = + torch::stable::detail::to(torch::stable::detail::from(dtype)); + auto shim_device_type = torch::stable::detail::to( + torch::stable::detail::from(device.type())); + auto shim_layout = + torch::stable::detail::to(torch::stable::detail::from(layout)); + AtenTensorHandle ath; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(aoti_torch_create_tensor_from_blob_v2( + data, + sizes.size(), + sizes.data(), + strides.data(), + storage_offset, + shim_dtype, + shim_device_type, + device.index(), + &ath, + shim_layout, + nullptr, + 0)); + return torch::stable::Tensor(ath); +} + +#if TORCH_FEATURE_VERSION >= (((0ULL + 2) << 56) | ((0ULL + 11) << 48)) +/// Creates a tensor from an existing data blob with a custom deleter. +/// +/// This is the same as the from_blob function above, but allows specifying a +/// custom deleter function that will be called when the tensor's storage is +/// deallocated. Accepts both plain function pointers and capturing lambdas. +/// +/// Minimum compatible version: PyTorch 2.11. +/// +/// @tparam F The callable type. Must be invocable with (void*). +/// @param data Pointer to the data buffer. +/// @param sizes The size of each dimension of the tensor. +/// @param strides The stride for each dimension. +/// @param device The device where the data resides. +/// @param dtype The scalar type of the data. +/// @param deleter Callable to invoke when the tensor is deallocated. +/// @param storage_offset The offset into the data buffer. Defaults to 0. +/// @param layout The memory layout. Defaults to Strided. +/// @return A tensor backed by the provided data. +template , int> = 0> +inline torch::stable::Tensor from_blob( + void* data, + torch::headeronly::IntHeaderOnlyArrayRef sizes, + torch::headeronly::IntHeaderOnlyArrayRef strides, + torch::stable::Device device, + torch::headeronly::ScalarType dtype, + F deleter, + int64_t storage_offset = 0, + torch::headeronly::Layout layout = torch::headeronly::Layout::Strided) { + auto shim_dtype = + torch::stable::detail::to(torch::stable::detail::from(dtype)); + auto shim_device_type = torch::stable::detail::to( + torch::stable::detail::from(device.type())); + auto shim_layout = + torch::stable::detail::to(torch::stable::detail::from(layout)); + + AtenTensorHandle ath; + if constexpr (std::is_convertible_v) { + // Simple function pointer: pass it as ctx, no heap allocation. + auto deleter_callback = [](void* data, void* ctx) { + auto fn = reinterpret_cast(ctx); + fn(data); + }; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_from_blob( + data, + sizes.size(), + sizes.data(), + strides.data(), + storage_offset, + shim_dtype, + shim_device_type, + device.index(), + &ath, + shim_layout, + nullptr, + 0, + deleter_callback, + reinterpret_cast(static_cast(deleter)))); + } else { + // Capturing lambda: heap-allocate and type-erase. + F* heap_allocated_deleter = new F(std::move(deleter)); + auto deleter_callback = [](void* data, void* ctx) { + F* func = static_cast(ctx); + (*func)(data); + delete func; + }; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_from_blob( + data, + sizes.size(), + sizes.data(), + strides.data(), + storage_offset, + shim_dtype, + shim_device_type, + device.index(), + &ath, + shim_layout, + nullptr, + 0, + deleter_callback, + static_cast(heap_allocated_deleter))); + } + return torch::stable::Tensor(ath); +} +#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_11_0 + +/// Stable version of the to.dtype_layout op. +/// +/// Converts a tensor to the specified dtype, layout, device, and/or memory +/// format. Returns a new tensor with the specified properties. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @param self The input tensor. +/// @param dtype Optional target scalar type. +/// @param layout Optional target memory layout. +/// @param device Optional target device. +/// @param pin_memory Optional flag to use pinned memory. +/// @param non_blocking If true, the operation may be asynchronous. Defaults to +/// false. +/// @param copy If true, always create a copy. Defaults to false. +/// @param memory_format Optional target memory format. +/// @return A tensor with the specified properties. +inline torch::stable::Tensor to( + const torch::stable::Tensor& self, + std::optional dtype = std::nullopt, + std::optional layout = std::nullopt, + std::optional device = std::nullopt, + std::optional pin_memory = std::nullopt, + bool non_blocking = false, + bool copy = false, + std::optional memory_format = + std::nullopt) { + const auto num_args = 8; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(dtype), + torch::stable::detail::from(layout), + torch::stable::detail::from(device), + torch::stable::detail::from(pin_memory), + torch::stable::detail::from(non_blocking), + torch::stable::detail::from(copy), + torch::stable::detail::from(memory_format)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::to", "dtype_layout", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Convenience overload for moving a tensor to a device. +/// +/// Moves the tensor to the specified device. This is a convenience wrapper +/// around the full to() function. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @param self The input tensor. +/// @param device The target device. +/// @param non_blocking If true, the operation may be asynchronous. Defaults to +/// false. +/// @param copy If true, always create a copy. Defaults to false. +/// @return A tensor on the specified device. +inline torch::stable::Tensor to( + const torch::stable::Tensor& self, + torch::stable::Device device, + bool non_blocking = false, + bool copy = false) { + return torch_stable_abi_shim::to( + self, + std::nullopt, + std::nullopt, + device, + std::nullopt, + non_blocking, + copy, + std::nullopt); +} + +/// Stable version of the contiguous op. +/// +/// Returns a contiguous in memory tensor containing the same data as the input +/// tensor. If the input tensor is already contiguous in the specified memory +/// format, the input tensor is returned. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @param self The input tensor. +/// @param memory_format The desired memory format. +/// @return A contiguous tensor. +inline torch::stable::Tensor contiguous( + const torch::stable::Tensor& self, + torch::headeronly::MemoryFormat memory_format = + torch::headeronly::MemoryFormat::Contiguous) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(memory_format)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::contiguous", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the new_empty op (2.10 version with full kwargs). +/// +/// Creates a new uninitialized tensor with the specified size and options. +/// This version supports all tensor creation kwargs. For versions < 2.10, +/// a simpler overload that only takes dtype is available. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @param self The input tensor whose properties may be inherited if kwargs are +/// not provided. +/// @param size The desired size of the output tensor. +/// @param dtype Optional scalar type for the tensor elements. +/// @param layout Optional memory layout (e.g., strided, sparse). +/// @param device Optional device to place the tensor on. +/// @param pin_memory Optional flag to use pinned memory (for CUDA tensors). +/// @return A new uninitialized tensor with the specified properties. +inline torch::stable::Tensor new_empty( + const torch::stable::Tensor& self, + torch::headeronly::IntHeaderOnlyArrayRef size, + std::optional dtype = std::nullopt, + std::optional layout = std::nullopt, + std::optional device = std::nullopt, + std::optional pin_memory = std::nullopt) { + const auto num_args = 6; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(size), + torch::stable::detail::from(dtype), + torch::stable::detail::from(layout), + torch::stable::detail::from(device), + torch::stable::detail::from(pin_memory)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::new_empty", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the new_zeros op (2.10 version with full kwargs). +/// +/// Creates a new zero-filled tensor with the specified size and options. +/// This version supports all tensor creation kwargs. For versions < 2.10, +/// a simpler overload that only takes dtype is available. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @param self The input tensor whose properties may be inherited if kwargs are +/// not provided. +/// @param size The desired size of the output tensor. +/// @param dtype Optional scalar type for the tensor elements. +/// @param layout Optional memory layout (e.g., strided, sparse). +/// @param device Optional device to place the tensor on. +/// @param pin_memory Optional flag to use pinned memory (for CUDA tensors). +/// @return A new zero-filled tensor with the specified properties. +inline torch::stable::Tensor new_zeros( + const torch::stable::Tensor& self, + torch::headeronly::IntHeaderOnlyArrayRef size, + std::optional dtype = std::nullopt, + std::optional layout = std::nullopt, + std::optional device = std::nullopt, + std::optional pin_memory = std::nullopt) { + const auto num_args = 6; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(size), + torch::stable::detail::from(dtype), + torch::stable::detail::from(layout), + torch::stable::detail::from(device), + torch::stable::detail::from(pin_memory)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::new_zeros", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the sum.dim_IntList op. +/// +/// Computes the sum of the input tensor along the specified dimensions. +/// If dim is not provided, sums over all dimensions. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @param self The input tensor. +/// @param dim Optional dimensions to reduce. If not provided, reduces all +/// dimensions. +/// @param keepdim Whether to retain the reduced dimensions. Defaults to false. +/// @param dtype Optional output dtype. If not provided, uses the input dtype. +/// @return A tensor containing the sum. +inline torch::stable::Tensor sum( + const torch::stable::Tensor& self, + std::optional dim = std::nullopt, + bool keepdim = false, + std::optional dtype = std::nullopt) { + const auto num_args = 4; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(dim), + torch::stable::detail::from(keepdim), + torch::stable::detail::from(dtype)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::sum", "dim_IntList", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the sum.IntList_out op. +/// +/// Computes the sum of the input tensor along the specified dimensions, +/// storing the result in the provided output tensor. Following C++ convention, +/// the out parameter comes first. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @param out The output tensor (modified in-place). +/// @param self The input tensor. +/// @param dim Optional dimensions to reduce. +/// @param keepdim Whether to retain the reduced dimensions. Defaults to false. +/// @param dtype Optional output dtype. +/// @return Reference to the output tensor. +inline torch::stable::Tensor& sum_out( + torch::stable::Tensor& out, + const torch::stable::Tensor& self, + std::optional dim = std::nullopt, + bool keepdim = false, + std::optional dtype = std::nullopt) { + const auto num_args = 5; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(dim), + torch::stable::detail::from(keepdim), + torch::stable::detail::from(dtype), + torch::stable::detail::from(out)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::sum", "IntList_out", stack.data(), TORCH_ABI_VERSION)); + // Clean up the handle in stack[0], discard the temporary + (void)torch::stable::detail::to(stack[0]); + return out; +} + +/// Stable version of the subtract.Tensor op. +/// +/// Subtracts the other tensor from self, with an optional scaling factor alpha. +/// Computes: self - alpha * other. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @note The alpha parameter is typed as double +/// API uses double for the Scalar parameter. +/// +/// @param self The input tensor. +/// @param other The tensor to subtract. +/// @param alpha The scaling factor for other. Defaults to 1.0. +/// @return The result of self - alpha * other. +inline torch::stable::Tensor subtract( + const torch::stable::Tensor& self, + const torch::stable::Tensor& other, + double alpha = 1.0) { + AtenTensorHandle ret0; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_aten_subtract_Tensor(self.get(), other.get(), alpha, &ret0)); + return torch::stable::Tensor(ret0); +} + +/// Stable version of the full.default op. +/// +/// Creates a tensor of the specified size filled with the given value. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @note The fill_value parameter is typed +/// C shim API uses double for the Scalar parameter. +/// +/// @param size The desired size of the output tensor. +/// @param fill_value The value to fill the tensor with. +/// @param dtype Optional scalar type for the tensor elements. +/// @param layout Optional memory layout. +/// @param device Optional device to place the tensor on. +/// @param pin_memory Optional flag to use pinned memory. +/// @return A new tensor filled with the specified value. +inline torch::stable::Tensor full( + torch::headeronly::IntHeaderOnlyArrayRef size, + double fill_value, + std::optional dtype = std::nullopt, + std::optional layout = std::nullopt, + std::optional device = std::nullopt, + std::optional pin_memory = std::nullopt) { + int32_t* dtype_ptr = nullptr; + int32_t dtype_val; + if (dtype.has_value()) { + dtype_val = torch::stable::detail::to( + torch::stable::detail::from(dtype.value())); + dtype_ptr = &dtype_val; + } + + int32_t* layout_ptr = nullptr; + int32_t layout_val; + if (layout.has_value()) { + layout_val = torch::stable::detail::to( + torch::stable::detail::from(layout.value())); + layout_ptr = &layout_val; + } + + int32_t* device_type_ptr = nullptr; + int32_t device_type_val; + int32_t device_index = 0; + if (device.has_value()) { + device_type_val = torch::stable::detail::to( + torch::stable::detail::from(device.value().type())); + device_type_ptr = &device_type_val; + device_index = device.value().index(); + } + + int32_t* pin_memory_ptr = nullptr; + int32_t pin_memory_val; + if (pin_memory.has_value()) { + pin_memory_val = pin_memory.value() ? 1 : 0; + pin_memory_ptr = &pin_memory_val; + } + + AtenTensorHandle ret0; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(aoti_torch_aten_full( + size.data(), + static_cast(size.size()), + fill_value, + dtype_ptr, + layout_ptr, + device_type_ptr, + device_index, + pin_memory_ptr, + &ret0)); + + return torch::stable::Tensor(ret0); +} + +/// Stable version of the bitwise_and.Tensor op. +/// +/// Computes the element-wise bitwise AND of self and other. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.14. +/// +/// @param self The first input tensor. +/// @param other The second input tensor. +/// @return The result of the bitwise AND. +inline torch::stable::Tensor bitwise_and( + const torch::stable::Tensor& self, + const torch::stable::Tensor& other) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(other)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::bitwise_and", "Tensor", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the bitwise_or.Tensor op. +/// +/// Computes the element-wise bitwise OR of self and other. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.14. +/// +/// @param self The first input tensor. +/// @param other The second input tensor. +/// @return The result of the bitwise OR. +inline torch::stable::Tensor bitwise_or( + const torch::stable::Tensor& self, + const torch::stable::Tensor& other) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(other)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::bitwise_or", "Tensor", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the bitwise_left_shift.Tensor op. +/// +/// Computes the element-wise left bitwise shift of self by other. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.14. +/// +/// @param self The input tensor to shift. +/// @param other The tensor of shift amounts. +/// @return The result of the left bitwise shift. +inline torch::stable::Tensor bitwise_left_shift( + const torch::stable::Tensor& self, + const torch::stable::Tensor& other) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(other)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::bitwise_left_shift", "Tensor", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the bitwise_right_shift.Tensor op. +/// +/// Computes the element-wise right bitwise shift of self by other. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.14. +/// +/// @param self The input tensor to shift. +/// @param other The tensor of shift amounts. +/// @return The result of the right bitwise shift. +inline torch::stable::Tensor bitwise_right_shift( + const torch::stable::Tensor& self, + const torch::stable::Tensor& other) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(other)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::bitwise_right_shift", "Tensor", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the permute op. +/// +/// Returns a view of the input tensor with its dimensions rearranged according +/// to dims. The returned tensor shares the same underlying data with the input. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.14. +/// +/// @param self The input tensor. +/// @param dims The desired ordering of dimensions. +/// @return A permuted view of the input tensor. +inline torch::stable::Tensor permute( + const torch::stable::Tensor& self, + torch::headeronly::IntHeaderOnlyArrayRef dims) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(dims)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::permute", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the index_select op. +/// +/// Returns a new tensor which indexes the input tensor along dimension dim +/// using the entries in index. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.14. +/// +/// @param self The input tensor. +/// @param dim The dimension to index along. +/// @param index A 1-D tensor of indices to select. +/// @return A tensor with values selected along dim. +inline torch::stable::Tensor index_select( + const torch::stable::Tensor& self, + int64_t dim, + const torch::stable::Tensor& index) { + const auto num_args = 3; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(dim), + torch::stable::detail::from(index)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::index_select", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the floor_divide op. +/// +/// Computes the element-wise floor division of self by other. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.14. +/// +/// @param self The dividend tensor. +/// @param other The divisor tensor. +/// @return The result of floor division. +inline torch::stable::Tensor floor_divide( + const torch::stable::Tensor& self, + const torch::stable::Tensor& other) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(other)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::floor_divide", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the is_pinned op. +/// +/// Returns true if the tensor's storage is allocated in pinned memory. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.14. +/// +/// @param self The input tensor. +/// @return true if the tensor is pinned, false otherwise. +inline bool is_pinned(const torch::stable::Tensor& self) { + const auto num_args = 2; + // Pass nullopt for aten::is_pinned's deprecated optional Device arg. + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(std::nullopt)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::is_pinned", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the sort op. +/// +/// Sorts the input tensor along dim, returning the values and their indices. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.15. +/// +/// @param self The input tensor. +/// @param dim The dimension to sort along. +/// @param descending Whether to sort in descending order. +/// @return A tuple of the sorted values and their indices. +inline std::tuple sort( + const torch::stable::Tensor& self, + int64_t dim = -1, + bool descending = false) { + const auto num_args = 3; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(dim), + torch::stable::detail::from(descending)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + torch_call_dispatcher("aten::sort", "", stack.data(), TORCH_ABI_VERSION)); + return std::make_tuple( + torch::stable::detail::to(stack[0]), + torch::stable::detail::to(stack[1])); +} + +/// Stable version of the sort.stable op. +/// +/// Sorts the input tensor along dim, returning the values and their indices. +/// This is a different overload than the sort op without the stable argument. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.15. +/// +/// @param self The input tensor. +/// @param stable Whether equal elements keep their original order. +/// @param dim The dimension to sort along. +/// @param descending Whether to sort in descending order. +/// @return A tuple of the sorted values and their indices. +inline std::tuple sort( + const torch::stable::Tensor& self, + std::optional stable, + int64_t dim = -1, + bool descending = false) { + const auto num_args = 4; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(stable), + torch::stable::detail::from(dim), + torch::stable::detail::from(descending)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::sort", "stable", stack.data(), TORCH_ABI_VERSION)); + return std::make_tuple( + torch::stable::detail::to(stack[0]), + torch::stable::detail::to(stack[1])); +} + +/// Stable version of the masked_select op. +/// +/// Returns a 1-D tensor of the elements of self where mask is true. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.15. +/// +/// @param self The input tensor. +/// @param mask A boolean tensor broadcastable with self. +/// @return A 1-D tensor of the selected elements. +inline torch::stable::Tensor masked_select( + const torch::stable::Tensor& self, + const torch::stable::Tensor& mask) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(mask)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::masked_select", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the mm.out op. +/// +/// Computes the matrix product of the 2-D tensors self and mat2, storing the +/// result in out. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.15. +/// +/// @param out The output tensor (modified in-place). +/// @param self The first 2-D input tensor. +/// @param mat2 The second 2-D input tensor. +/// @return Reference to the output tensor. +inline torch::stable::Tensor& mm_out( + torch::stable::Tensor& out, + const torch::stable::Tensor& self, + const torch::stable::Tensor& mat2) { + const auto num_args = 3; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(mat2), + torch::stable::detail::from(out)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::mm", "out", stack.data(), TORCH_ABI_VERSION)); + // Clean up the handle in stack[0], discard the temporary + (void)torch::stable::detail::to(stack[0]); + return out; +} + +/// Stable version of the mm.dtype_out op. +/// +/// Computes the matrix product of self and mat2 in out_dtype. This is a +/// different overload than mm_out without the out_dtype argument. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.15. +/// +/// @param out The output tensor (modified in-place). +/// @param self The first 2-D input tensor. +/// @param mat2 The second 2-D input tensor. +/// @param out_dtype The dtype of the result. +/// @return Reference to the output tensor. +inline torch::stable::Tensor& mm_out( + torch::stable::Tensor& out, + const torch::stable::Tensor& self, + const torch::stable::Tensor& mat2, + torch::headeronly::ScalarType out_dtype) { + const auto num_args = 4; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(mat2), + torch::stable::detail::from(out_dtype), + torch::stable::detail::from(out)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::mm", "dtype_out", stack.data(), TORCH_ABI_VERSION)); + // Clean up the handle in stack[0], discard the temporary + (void)torch::stable::detail::to(stack[0]); + return out; +} + +/// Stable version of the bmm.out op. +/// +/// Computes the batch matrix product of the 3-D tensors self and mat2, +/// storing the result in out. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.15. +/// +/// @param out The output tensor (modified in-place). +/// @param self The first 3-D input tensor of shape (b, n, m). +/// @param mat2 The second 3-D input tensor of shape (b, m, p). +/// @return Reference to the output tensor. +inline torch::stable::Tensor& bmm_out( + torch::stable::Tensor& out, + const torch::stable::Tensor& self, + const torch::stable::Tensor& mat2) { + const auto num_args = 3; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(mat2), + torch::stable::detail::from(out)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::bmm", "out", stack.data(), TORCH_ABI_VERSION)); + // Clean up the handle in stack[0], discard the temporary + (void)torch::stable::detail::to(stack[0]); + return out; +} + +/// Stable version of the bmm.dtype_out op. +/// +/// Computes the batch matrix product of self and mat2 in out_dtype. This is a +/// different overload than bmm_out without the out_dtype argument. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.15. +/// +/// @param out The output tensor (modified in-place). +/// @param self The first 3-D input tensor of shape (b, n, m). +/// @param mat2 The second 3-D input tensor of shape (b, m, p). +/// @param out_dtype The dtype of the result. +/// @return Reference to the output tensor. +inline torch::stable::Tensor& bmm_out( + torch::stable::Tensor& out, + const torch::stable::Tensor& self, + const torch::stable::Tensor& mat2, + torch::headeronly::ScalarType out_dtype) { + const auto num_args = 4; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(mat2), + torch::stable::detail::from(out_dtype), + torch::stable::detail::from(out)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::bmm", "dtype_out", stack.data(), TORCH_ABI_VERSION)); + // Clean up the handle in stack[0], discard the temporary + (void)torch::stable::detail::to(stack[0]); + return out; +} + +/// Stable version of the flip op. +/// +/// Reverses the order of elements along the given dimensions and returns +/// the result as a new tensor. Negative dimensions are wrapped. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.15. +/// +/// @param self The input tensor. +/// @param dims The dimensions to flip. +/// @return A new tensor with the given dimensions reversed. +inline torch::stable::Tensor flip( + const torch::stable::Tensor& self, + torch::headeronly::IntHeaderOnlyArrayRef dims) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(dims)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + torch_call_dispatcher("aten::flip", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the from_file op. +/// +/// Creates a CPU tensor backed by a memory-map of the file at filename. +/// Unlike at::from_file, filename is a const std::string& rather than a +/// std::string_view. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.15. +/// +/// @param filename Path of the file to read. +/// @param shared Optional flag to write tensor changes back to the file. +/// @param size Optional number of elements to read from the file. +/// @param dtype Optional scalar type for the tensor elements. +/// @param layout Optional memory layout, which must be strided for this op. +/// @param device Optional device, which must be CPU for this op. +/// @param pin_memory Optional pinning flag, which must be false for this op. +/// @return A new tensor holding the file contents. +inline torch::stable::Tensor from_file( + const std::string& filename, + std::optional shared = std::nullopt, + std::optional size = 0, + std::optional dtype = std::nullopt, + std::optional layout = std::nullopt, + std::optional device = std::nullopt, + std::optional pin_memory = std::nullopt) { + const auto num_args = 7; + std::array stack{ + torch::stable::detail::from(filename), + torch::stable::detail::from(shared), + torch::stable::detail::from(size), + torch::stable::detail::from(dtype), + torch::stable::detail::from(layout), + torch::stable::detail::from(device), + torch::stable::detail::from(pin_memory)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::from_file", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + +#if TORCH_FEATURE_VERSION >= (((0ULL + 2) << 56) | ((0ULL + 16) << 48)) + +/// Stable version of the add.Tensor op. +/// +/// Adds the other tensor to self, scaled by alpha: self + alpha * other. +/// Unlike at::add, alpha is a double rather than a Scalar. Adding two +/// integral or boolean tensors is not supported. +/// +/// Minimum compatible version: PyTorch 2.16. +/// Upstream header availability: PyTorch 2.16. +/// +/// @param self The input tensor. +/// @param other The tensor to add. +/// @param alpha The scaling factor for other. Defaults to 1.0. +/// @return The result of self + alpha * other. +inline torch::stable::Tensor add( + const torch::stable::Tensor& self, + const torch::stable::Tensor& other, + double alpha = 1.0) { + AtenTensorHandle ret0; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_aten_add_Tensor(self.get(), other.get(), alpha, &ret0)); + return torch::stable::Tensor(ret0); +} + +#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_16_0 + +} // namespace torch_stable_abi_shim