chore: import upstream snapshot with attribution

This commit is contained in:
wehub-resource-sync
2026-07-13 12:40:42 +08:00
commit e25996e7db
15472 changed files with 3536181 additions and 0 deletions
+37
View File
@@ -0,0 +1,37 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/Device.h>
#include <ATen/Functions.h>
#include <ATen/Tensor.h>
#include <ATen/TensorIndexing.h>
#include <ATen/Utils.h>
#include <c10/core/Device.h>
#include <c10/core/DeviceType.h>
#include <c10/core/MemoryFormat.h>
#include <c10/core/Scalar.h>
#include <c10/core/ScalarType.h>
#include <c10/core/SymIntArrayRef.h>
#include <c10/core/TensorOptions.h>
#include <c10/util/ArrayRef.h>
#include <c10/util/Exception.h>
#include <c10/util/OptionalArrayRef.h>
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
#include <c10/cuda/CUDAException.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#endif
@@ -0,0 +1,49 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#include <ATen/AccumulateType.h>
namespace at {
c10::ScalarType toAccumulateType(c10::ScalarType type, c10::DeviceType device) {
switch (type) {
#define DEFINE_CASE(scalar_t, TypeNum) \
case ScalarType::TypeNum: \
switch (device) { \
case DeviceType::CUDA: \
return CppTypeToScalarType< \
at::acc_type_device<scalar_t, c10::DeviceType::CUDA>>::value; \
default: \
return CppTypeToScalarType< \
at::acc_type_device<scalar_t, c10::DeviceType::CPU>>::value; \
}
AT_FORALL_SCALAR_TYPES_WITH_COMPLEX_EXCEPT_COMPLEX_HALF_F8NZ(DEFINE_CASE)
#undef DEFINE_CASE
default:
TORCH_INTERNAL_ASSERT(false, "Unrecognized ScalarType: ", type);
}
}
c10::ScalarType toAccumulateType(c10::ScalarType type, bool is_cuda) {
return is_cuda ? toAccumulateType(type, c10::DeviceType::CUDA)
: toAccumulateType(type, c10::DeviceType::CPU);
}
} // namespace at
@@ -0,0 +1,115 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#pragma once
#include <c10/core/DeviceType.h>
#include <c10/core/ScalarType.h>
#include <c10/util/BFloat16.h>
#include <c10/util/Float8_e4m3fn.h>
// #include <c10/util/Float8_e4m3fnuz.h>
#include <c10/util/Float8_e5m2.h>
// #include <c10/util/Float8_e5m2fnuz.h>
#include <c10/util/Half.h>
#if defined(__CUDACC__)
#include <cuda.h>
#include <cuda_fp16.h>
#elif defined(__HIPCC__)
#include <hip/hip_fp16.h>
#include <hip/hip_runtime.h>
#endif
namespace at {
template <typename T, c10::DeviceType D>
struct AccumulateTypeDevice {};
template <typename T, bool>
struct AccumulateType {};
template <typename T>
struct AccumulateType<T, false> {
using type = typename AccumulateTypeDevice<T, c10::DeviceType::CPU>::type;
};
template <typename T>
struct AccumulateType<T, true> {
using type = typename AccumulateTypeDevice<T, c10::DeviceType::CUDA>::type;
};
template <typename T, c10::DeviceType device>
using acc_type_device = typename AccumulateTypeDevice<T, device>::type;
template <typename T, bool is_cuda>
using acc_type = typename AccumulateType<T, is_cuda>::type;
#define ACC_TYPE(t, acc_t, device_type) \
template <> \
struct AccumulateTypeDevice<t, device_type> { \
using type = acc_t; \
};
#define CUDA_ACC_TYPE(t, acc_t) ACC_TYPE(t, acc_t, c10::DeviceType::CUDA)
#define CPU_ACC_TYPE(t, acc_t) ACC_TYPE(t, acc_t, c10::DeviceType::CPU)
#if defined(__CUDACC__) || defined(__HIPCC__)
CUDA_ACC_TYPE(half, float)
#endif
CUDA_ACC_TYPE(BFloat16, float)
CUDA_ACC_TYPE(Half, float)
CUDA_ACC_TYPE(Float8_e5m2, float)
CUDA_ACC_TYPE(Float8_e4m3fn, float)
// CUDA_ACC_TYPE(Float8_e5m2fnuz, float)
// CUDA_ACC_TYPE(Float8_e4m3fnuz, float)
CUDA_ACC_TYPE(float, float)
CUDA_ACC_TYPE(double, double)
CUDA_ACC_TYPE(int8_t, int64_t)
CUDA_ACC_TYPE(uint8_t, int64_t)
CUDA_ACC_TYPE(char, int64_t)
CUDA_ACC_TYPE(int16_t, int64_t)
CUDA_ACC_TYPE(int32_t, int64_t)
CUDA_ACC_TYPE(int64_t, int64_t)
CUDA_ACC_TYPE(bool, bool)
CUDA_ACC_TYPE(c10::complex<Half>, c10::complex<float>)
CUDA_ACC_TYPE(c10::complex<float>, c10::complex<float>)
CUDA_ACC_TYPE(c10::complex<double>, c10::complex<double>)
CPU_ACC_TYPE(BFloat16, float)
CPU_ACC_TYPE(Half, float)
CPU_ACC_TYPE(Float8_e5m2, float)
CPU_ACC_TYPE(Float8_e4m3fn, float)
// CPU_ACC_TYPE(Float8_e5m2fnuz, float)
// CPU_ACC_TYPE(Float8_e4m3fnuz, float)
CPU_ACC_TYPE(float, double)
CPU_ACC_TYPE(double, double)
CPU_ACC_TYPE(int8_t, int64_t)
CPU_ACC_TYPE(uint8_t, int64_t)
CPU_ACC_TYPE(char, int64_t)
CPU_ACC_TYPE(int16_t, int64_t)
CPU_ACC_TYPE(int32_t, int64_t)
CPU_ACC_TYPE(int64_t, int64_t)
CPU_ACC_TYPE(bool, bool)
CPU_ACC_TYPE(c10::complex<Half>, c10::complex<float>)
CPU_ACC_TYPE(c10::complex<float>, c10::complex<double>)
CPU_ACC_TYPE(c10::complex<double>, c10::complex<double>)
c10::ScalarType toAccumulateType(c10::ScalarType type, c10::DeviceType device);
c10::ScalarType toAccumulateType(c10::ScalarType type, bool is_cuda);
} // namespace at
@@ -0,0 +1,16 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <c10/core/Device.h>
@@ -0,0 +1,35 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/Tensor.h>
#include <c10/core/ScalarType.h>
#include <c10/util/Optional.h>
namespace at {
inline std::optional<Device> device_of(const Tensor& t) {
if (t.defined()) {
return t.device();
} else {
return std::nullopt;
}
}
inline std::optional<Device> device_of(const std::optional<Tensor>& t) {
return t.has_value() ? device_of(t.value()) : std::nullopt;
}
} // namespace at
@@ -0,0 +1,82 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/ops/_local_scalar_dense.h>
#include <ATen/ops/_nnz.h>
#include <ATen/ops/_values.h>
#include <ATen/ops/abs.h>
#include <ATen/ops/all.h>
#include <ATen/ops/allclose.h>
#include <ATen/ops/any.h>
#include <ATen/ops/arange.h>
#include <ATen/ops/as_strided.h>
#include <ATen/ops/cat.h>
#include <ATen/ops/chunk.h>
#include <ATen/ops/clamp.h>
#include <ATen/ops/coalesce.h>
#include <ATen/ops/detach.h>
#include <ATen/ops/dsplit.h>
#include <ATen/ops/empty.h>
#include <ATen/ops/empty_like.h>
#include <ATen/ops/empty_strided.h>
#include <ATen/ops/equal.h>
#include <ATen/ops/expand.h>
#include <ATen/ops/expand_as.h>
#include <ATen/ops/eye.h>
#include <ATen/ops/flatten.h>
#include <ATen/ops/from_blob.h>
#include <ATen/ops/full.h>
#include <ATen/ops/hsplit.h>
#include <ATen/ops/index.h>
#include <ATen/ops/index_put.h>
#include <ATen/ops/is_coalesced.h>
#include <ATen/ops/item.h>
#include <ATen/ops/masked_select.h>
#include <ATen/ops/narrow.h>
#include <ATen/ops/narrow_copy.h>
#include <ATen/ops/new_empty.h>
#include <ATen/ops/new_full.h>
#include <ATen/ops/new_ones.h>
#include <ATen/ops/new_zeros.h>
#include <ATen/ops/ones.h>
#include <ATen/ops/permute.h>
#include <ATen/ops/reciprocal.h>
#include <ATen/ops/record_stream.h>
#include <ATen/ops/rename.h>
#include <ATen/ops/reshape.h>
#include <ATen/ops/resize.h>
#include <ATen/ops/select.h>
#include <ATen/ops/slice.h>
#include <ATen/ops/sparse_coo_tensor.h>
#include <ATen/ops/sparse_csr_tensor.h>
#include <ATen/ops/split.h>
#include <ATen/ops/split_with_sizes.h>
#include <ATen/ops/squeeze.h>
#include <ATen/ops/std.h>
#include <ATen/ops/sum.h>
#include <ATen/ops/t.h>
#include <ATen/ops/tensor_split.h>
#include <ATen/ops/to.h>
#include <ATen/ops/transpose.h>
#include <ATen/ops/unflatten.h>
#include <ATen/ops/unsafe_split.h>
#include <ATen/ops/unsafe_split_with_sizes.h>
#include <ATen/ops/unsqueeze.h>
#include <ATen/ops/view.h>
#include <ATen/ops/view_as.h>
#include <ATen/ops/vsplit.h>
#include <ATen/ops/zeros.h>
#include <ATen/ops/zeros_like.h>
@@ -0,0 +1,69 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <c10/core/ScalarType.h>
#include <c10/util/BFloat16.h>
#include <c10/util/Exception.h>
#include <c10/util/Float8_e4m3fn.h>
#include <c10/util/Float8_e5m2.h>
#include <c10/util/Half.h>
namespace at {
// For FP16 or BFloat16 inputs, ops should perform internal math in FP32.
template <typename scalar_t>
struct OpMathType {
using type = scalar_t;
};
template <>
struct OpMathType<at::Half> {
using type = float;
};
template <>
struct OpMathType<at::BFloat16> {
using type = float;
};
template <>
struct OpMathType<at::Float8_e5m2> {
using type = float;
};
template <>
struct OpMathType<at::Float8_e4m3fn> {
using type = float;
};
template <>
struct OpMathType<c10::complex<Half>> {
using type = c10::complex<float>;
};
template <typename T>
using opmath_type = typename OpMathType<T>::type;
inline c10::ScalarType toOpMathType(const c10::ScalarType type) {
switch (type) {
#define DEFINE_CASE(scalar_t, TypeNum) \
case ScalarType::TypeNum: \
return CppTypeToScalarType<at::opmath_type<scalar_t>>::value;
AT_FORALL_SCALAR_TYPES_WITH_COMPLEX(DEFINE_CASE)
#undef DEFINE_CASE
default:
TORCH_INTERNAL_ASSERT(false, "Unrecognized ScalarType: ", type);
}
}
} // namespace at
@@ -0,0 +1,17 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
@@ -0,0 +1,124 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <c10/core/SymInt.h>
#include <cstdint>
#include <cstring>
#include <limits>
#include <memory>
#include <optional>
#include <stdexcept>
#include <type_traits>
#include <utility>
namespace at {
class Tensor;
}
namespace at::indexing {
constexpr int64_t INDEX_MIN = std::numeric_limits<int64_t>::min();
constexpr int64_t INDEX_MAX = std::numeric_limits<int64_t>::max();
enum class TensorIndexType { None, Ellipsis, SymInt, Boolean, Slice, Tensor };
constexpr std::nullopt_t None = std::nullopt;
struct EllipsisIndexType final {
EllipsisIndexType() = default;
};
const EllipsisIndexType Ellipsis = EllipsisIndexType();
struct Slice final {
public:
Slice(std::optional<c10::SymInt> start_index = std::nullopt,
std::optional<c10::SymInt> stop_index = std::nullopt,
std::optional<c10::SymInt> step_index = std::nullopt) {
if (!step_index.has_value()) {
step_ = c10::SymInt(1);
} else {
step_ = std::move(step_index).value();
}
if (!start_index.has_value()) {
start_ = c10::SymInt(step_ < 0 ? INDEX_MAX : 0);
} else {
start_ = std::move(start_index).value();
}
if (!stop_index.has_value()) {
stop_ = c10::SymInt(step_ < 0 ? INDEX_MIN : INDEX_MAX);
} else {
stop_ = std::move(stop_index).value();
}
}
inline c10::SymInt start() const { return start_; }
inline c10::SymInt stop() const { return stop_; }
inline c10::SymInt step() const { return step_; }
private:
c10::SymInt start_;
c10::SymInt stop_;
c10::SymInt step_;
};
struct TensorIndex final {
TensorIndex(std::nullopt_t /*unused*/) : type_(TensorIndexType::None) {}
TensorIndex(at::indexing::EllipsisIndexType /*unused*/)
: type_(TensorIndexType::Ellipsis) {}
TensorIndex(const char* str) : TensorIndex(at::indexing::Ellipsis) {
if (std::strcmp(str, "...") != 0) {
throw std::invalid_argument(
"Expected \"...\" to represent an ellipsis index.");
}
}
TensorIndex(c10::SymInt integer)
: integer_(std::move(integer)), type_(TensorIndexType::SymInt) {}
TensorIndex(int64_t integer) : TensorIndex(c10::SymInt(integer)) {}
TensorIndex(int integer) : TensorIndex(c10::SymInt(integer)) {}
template <class T, class = std::enable_if_t<std::is_same_v<bool, T>>>
TensorIndex(T boolean) : boolean_(boolean), type_(TensorIndexType::Boolean) {}
TensorIndex(Slice slice)
: slice_(std::move(slice)), type_(TensorIndexType::Slice) {}
TensorIndex(const at::Tensor& tensor);
inline bool is_none() const { return type_ == TensorIndexType::None; }
inline bool is_ellipsis() const { return type_ == TensorIndexType::Ellipsis; }
inline bool is_integer() const { return type_ == TensorIndexType::SymInt; }
inline c10::SymInt integer() const { return integer_; }
inline bool is_boolean() const { return type_ == TensorIndexType::Boolean; }
inline bool boolean() const { return boolean_; }
inline bool is_slice() const { return type_ == TensorIndexType::Slice; }
inline const Slice& slice() const { return slice_; }
inline bool is_tensor() const { return type_ == TensorIndexType::Tensor; }
const at::Tensor& tensor() const;
private:
c10::SymInt integer_ = 0;
bool boolean_ = false;
Slice slice_;
std::shared_ptr<at::Tensor> tensor_;
TensorIndexType type_;
};
} // namespace at::indexing
@@ -0,0 +1,97 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <ATen/Utils.h>
#include <ATen/ops/empty.h>
#include <ATen/ops/to.h>
#include <c10/core/Layout.h>
#include <c10/core/ScalarType.h>
#include <c10/util/ArrayRef.h>
#include <c10/util/Exception.h>
#include <c10/util/accumulate.h>
#include <algorithm>
#include "paddle/common/macros.h"
#include "paddle/phi/api/include/sparse_api.h"
#include "paddle/phi/api/include/tensor.h"
namespace at {
namespace detail {
template <typename T>
Tensor tensor_cpu(ArrayRef<T> values, const TensorOptions& options) {
constexpr auto native_scalar_type = c10::CppTypeToScalarType<T>::value;
auto result = at::empty(values.size(), options.dtype(native_scalar_type));
PD_CHECK(result.is_contiguous());
std::copy(values.begin(), values.end(), result.template data_ptr<T>());
if (options.dtype() != native_scalar_type) {
return result.to(at::TensorOptions().dtype(options.dtype()));
}
return result;
}
template <typename T>
Tensor tensor_backend(ArrayRef<T> values, const TensorOptions& options) {
auto cpu_tensor =
tensor_cpu(values, options.device(c10::Device(c10::DeviceType::CPU)));
return cpu_tensor.to(options.device());
}
template <typename T>
Tensor tensor_complex_cpu(ArrayRef<T> values, const TensorOptions& options) {
constexpr auto native_scalar_type = c10::CppTypeToScalarType<T>::value;
auto result = at::empty(values.size(), options.dtype(native_scalar_type));
PD_CHECK(result.is_contiguous());
std::copy(values.begin(), values.end(), result.template data_ptr<T>());
if (options.dtype() != native_scalar_type) {
return result.to(at::TensorOptions().dtype(options.dtype()));
}
return result;
}
template <typename T>
Tensor tensor_complex_backend(ArrayRef<T> values,
const TensorOptions& options) {
auto cpu_tensor = tensor_complex_cpu(
values, options.device(c10::Device(c10::DeviceType::CPU)));
return cpu_tensor.to(options.device());
}
} // namespace detail
#define TENSOR(T, _1) \
PADDLE_API Tensor tensor(ArrayRef<T> values, const TensorOptions& options) { \
if (options.device().type() != c10::DeviceType::CPU) { \
return at::detail::tensor_backend(values, options); \
} else { \
return at::detail::tensor_cpu(values, options); \
} \
}
AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, TENSOR)
#undef TENSOR
#define TENSOR(T, _1) \
PADDLE_API Tensor tensor(ArrayRef<T> values, const TensorOptions& options) { \
if (options.device().type() != c10::DeviceType::CPU) { \
return at::detail::tensor_complex_backend(values, options); \
} else { \
return at::detail::tensor_complex_cpu(values, options); \
} \
}
AT_FORALL_COMPLEX_TYPES(TENSOR)
#undef TENSOR
} // namespace at
@@ -0,0 +1,36 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
namespace at {
namespace detail {
template <typename T>
Tensor tensor_cpu(ArrayRef<T> values, const TensorOptions& options);
template <typename T>
Tensor tensor_backend(ArrayRef<T> values, const TensorOptions& options);
template <typename T>
Tensor tensor_complex_cpu(ArrayRef<T> values, const TensorOptions& options);
template <typename T>
Tensor tensor_complex_backend(ArrayRef<T> values, const TensorOptions& options);
} // namespace detail
} // namespace at
@@ -0,0 +1,197 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#pragma once
#include <c10/core/Device.h>
#include <c10/core/DispatchKeySet.h>
#include <c10/core/GeneratorImpl.h>
#include <c10/util/Exception.h>
#include <c10/util/intrusive_ptr.h>
#include <cstdint>
#include <deque>
#include <memory>
#include <mutex> // NOLINT(build/c++11)
#include <optional>
#include <utility>
/**
* Note [Generator]
* ~~~~~~~~~~~~~~~~
* A Pseudo Random Number Generator (PRNG) is an engine that uses an algorithm
* to generate a seemingly random sequence of numbers, that may be later be
* used in creating a random distribution. Such an engine almost always
* maintains a state and requires a seed to start off the creation of random
* numbers. Often times, users have found it beneficial to be able to
* explicitly create, retain, and destroy PRNG states and also be able to
* have control over the seed value.
*
* A Generator in ATen gives users the ability to read, write and modify a
* PRNG engine. For instance, it does so by letting users seed a PRNG engine,
* fork the state of the engine, etc.
*
* By default, there is one generator per device, and a device's generator is
* lazily created. A user can use the torch.Generator() api to create their
* own generator.
*
* This implementation wraps Paddle's phi::Generator via c10::GeneratorImpl.
*/
/**
* Note [Acquire lock when using random generators]
* ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
* Generator and its derived classes are NOT thread-safe. Please use the
* public mutex_ when using any methods from these classes, except for the
* read-only methods.
*/
// Forward declare PyObject if not already available.
#ifndef PyObject
struct _object;
using PyObject = _object;
#endif
namespace at {
using c10::Device;
using c10::DispatchKeySet;
class Tensor;
struct Generator {
Generator() = default;
explicit Generator(c10::intrusive_ptr<c10::GeneratorImpl> gen_impl)
: impl_(std::move(gen_impl)) {
TORCH_CHECK(impl_.get(), "GeneratorImpl with nullptr is not supported");
}
bool operator==(const Generator& rhs) const {
return this->impl_ == rhs.impl_;
}
bool operator!=(const Generator& rhs) const { return !((*this) == rhs); }
bool defined() const { return static_cast<bool>(impl_); }
c10::GeneratorImpl* unsafeGetGeneratorImpl() const { return impl_.get(); }
c10::GeneratorImpl* unsafeReleaseGeneratorImpl() { return impl_.release(); }
const c10::intrusive_ptr<c10::GeneratorImpl>& getIntrusivePtr() const {
return impl_;
}
void set_current_seed(uint64_t seed) { impl_->set_current_seed(seed); }
/// Sets the offset of Generator state to the desired offset.
/// Supported for Philox based Generators (CUDA / MPS).
void set_offset(uint64_t offset) { impl_->set_offset(offset); }
/// Returns the offset of Generator state.
/// Supported for Philox based Generators (CUDA / MPS).
uint64_t get_offset() const { return impl_->get_offset(); }
uint64_t current_seed() const { return impl_->current_seed(); }
uint64_t seed() { return impl_->seed(); }
// ----- state transfer (not inlined to break header cycles) ----------------
// These methods mirror PyTorch's set_state / get_state which operate on
// serialised byte tensors. In the Paddle compat layer we provide a simpler
// state-copy semantic through graphsafe_set_state / graphsafe_get_state.
/// Copy the full PRNG state from another Generator.
void graphsafe_set_state(const Generator& src) {
TORCH_CHECK(src.defined(), "Source generator is not defined");
TORCH_CHECK(defined(), "Target generator is not defined");
auto src_state = src.impl_->paddle_generator()->GetState();
impl_->paddle_generator()->SetState(src_state);
}
/// Obtain a Generator whose state is a snapshot (clone) of this one.
Generator graphsafe_get_state() const {
TORCH_CHECK(defined(), "Generator is not defined");
return clone();
}
std::mutex& mutex() { return impl_->mutex_; }
DispatchKeySet key_set() const { return impl_->key_set(); }
Device device() const { return impl_->device(); }
inline void set_pyobj(PyObject* pyobj) const noexcept {
impl_->set_pyobj(pyobj);
}
inline PyObject* pyobj() const noexcept { return impl_->pyobj(); }
template <typename T>
T* get() const {
return static_cast<T*>(impl_.get());
}
Generator clone() const { return Generator(impl_->clone()); }
/// Access the underlying Paddle phi::Generator (convenience).
std::shared_ptr<phi::Generator> paddle_generator() const {
return impl_->paddle_generator();
}
private:
c10::intrusive_ptr<c10::GeneratorImpl> impl_;
};
template <class Impl, class... Args>
Generator make_generator(Args&&... args) {
return Generator(c10::make_intrusive<Impl>(std::forward<Args>(args)...));
}
/**
* Utility function to static cast input Generator to
* the backend generator type (CPUGeneratorImpl / CUDAGeneratorImpl etc.)
*/
template <typename T>
inline T* check_generator(std::optional<Generator> gen) {
TORCH_CHECK(gen.has_value(), "Expected Generator but received nullopt");
TORCH_CHECK(gen->defined(),
"Generator with undefined implementation is not allowed");
TORCH_CHECK(
T::device_type() == gen->device().type(),
"Expected a generator for ",
phi::AllocationTypeStr(c10::DeviceTypeToPhi(T::device_type())),
" but found one for ",
phi::AllocationTypeStr(c10::DeviceTypeToPhi(gen->device().type())));
return gen->get<T>();
}
/**
* Utility function used in tensor implementations, which supplies the
* default generator to tensors if an input generator is not supplied.
* The input Generator is also static-cast to the backend generator type.
*/
template <typename T>
inline T* get_generator_or_default(const std::optional<Generator>& gen,
const Generator& default_gen) {
return gen.has_value() && gen->defined() ? check_generator<T>(gen)
: check_generator<T>(default_gen);
}
} // namespace at
@@ -0,0 +1,15 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <c10/core/Scalar.h>
@@ -0,0 +1,17 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/TensorBody.h>
@@ -0,0 +1,104 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#pragma once
#include <torch/headeronly/core/TensorAccessor.h>
#include <c10/macros/Macros.h>
#include <c10/util/ArrayRef.h>
#include <c10/util/Exception.h>
#include <cstddef>
#include <cstdint>
#include <type_traits>
namespace at {
using torch::headeronly::DefaultPtrTraits;
#if defined(__CUDACC__) || defined(__HIPCC__)
using torch::headeronly::RestrictPtrTraits;
#endif
template <typename T,
size_t N,
template <typename U> class PtrTraits = DefaultPtrTraits,
typename index_t = int64_t>
using TensorAccessorBase = torch::headeronly::detail::
TensorAccessorBase<c10::IntArrayRef, T, N, PtrTraits, index_t>;
template <typename T,
size_t N,
template <typename U> class PtrTraits = DefaultPtrTraits,
typename index_t = int64_t>
using TensorAccessor = torch::headeronly::detail::
TensorAccessor<c10::IntArrayRef, T, N, PtrTraits, index_t>;
namespace detail {
template <size_t N, typename index_t>
struct IndexBoundsCheck {
explicit IndexBoundsCheck(index_t i) {
TORCH_CHECK(0 <= i && i < index_t{N},
"Index ",
i,
" is not within bounds of a tensor of dimension ",
N);
}
};
} // namespace detail
template <typename T,
size_t N,
template <typename U> class PtrTraits = DefaultPtrTraits,
typename index_t = int64_t>
using GenericPackedTensorAccessorBase =
torch::headeronly::detail::GenericPackedTensorAccessorBase<
detail::IndexBoundsCheck<N, index_t>,
T,
N,
PtrTraits,
index_t>;
template <typename T,
size_t N,
template <typename U> class PtrTraits = DefaultPtrTraits,
typename index_t = int64_t>
using GenericPackedTensorAccessor =
torch::headeronly::detail::GenericPackedTensorAccessor<
TensorAccessor<T, N - 1, PtrTraits, index_t>,
detail::IndexBoundsCheck<N, index_t>,
T,
N,
PtrTraits,
index_t>;
template <typename T,
size_t N,
template <typename U> class PtrTraits = DefaultPtrTraits>
using PackedTensorAccessor32 =
GenericPackedTensorAccessor<T, N, PtrTraits, int32_t>;
template <typename T,
size_t N,
template <typename U> class PtrTraits = DefaultPtrTraits>
using PackedTensorAccessor64 =
GenericPackedTensorAccessor<T, N, PtrTraits, int64_t>;
} // namespace at
@@ -0,0 +1,578 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/TensorAccessor.h>
#include <c10/core/Device.h>
#include <c10/core/Layout.h>
#include <c10/core/MemoryFormat.h>
#include <c10/core/Scalar.h>
#include <c10/core/ScalarType.h>
#include <c10/core/Storage.h>
#include <c10/core/SymInt.h>
#include <c10/core/SymIntArrayRef.h>
#include <c10/core/TensorOptions.h>
#include <utils/int_array_ref_conversion.h>
#include <utils/scalar_type_conversion.h>
#include <algorithm>
#include <iostream>
#include <memory>
#include <mutex>
#include <type_traits>
#include <unordered_map>
#include <vector>
#include "paddle/common/layout.h"
#include "paddle/phi/api/include/api.h"
#include "paddle/phi/api/include/tensor.h"
#include "paddle/phi/common/place.h"
#include "paddle/phi/core/dense_tensor.h"
namespace at {
using PaddleTensor = paddle::Tensor;
class PADDLE_API TensorBase {
public:
TensorBase() = default;
explicit TensorBase(const PaddleTensor& tensor) : tensor_(tensor) {
InitStorage();
}
TensorBase(const TensorBase&) = default;
TensorBase(TensorBase&&) noexcept = default;
~TensorBase() noexcept = default;
#if defined(_MSC_VER)
TensorBase& operator=(const TensorBase& x) & {
tensor_ = x.tensor_;
storage_ = x.storage_;
return *this;
}
TensorBase& operator=(TensorBase&& x) & noexcept {
tensor_ = std::move(x.tensor_);
storage_ = std::move(x.storage_);
return *this;
}
#else
TensorBase& operator=(const TensorBase& x) & = default;
TensorBase& operator=(TensorBase&& x) & noexcept = default;
#endif
TensorBase& operator=(const TensorBase&) && = delete;
TensorBase& operator=(TensorBase&&) && noexcept = delete;
bool is_same(const TensorBase& other) const noexcept {
return tensor_.impl().get() == other.tensor_.impl().get();
}
size_t use_count() const noexcept { return tensor_.impl().use_count(); }
size_t weak_use_count() const noexcept {
// PyTorch exposes an internal self weak-reference on live TensorImpls, so
// the observable weak count starts at 1 even without user-created refs.
return tensor_.defined() ? 1 : 0;
}
void print() const {
if (defined()) {
std::cerr << '[' << toString() << ' ' << sizes() << ']' << '\n';
} else {
std::cerr << "[UndefinedTensor]" << '\n';
}
}
std::string toString() const {
if (!tensor_.defined()) {
return "UndefinedType";
}
std::string backend_str;
const auto& place = tensor_.place();
if (phi::is_cpu_place(place)) {
backend_str = "CPU";
} else if (phi::is_gpu_place(place)) {
backend_str = "CUDA";
} else {
backend_str = "Undefined";
}
std::string scalar_type_str = at::toString(scalar_type());
return backend_str + scalar_type_str + "Type";
}
// Returns the tensor's current data pointer. Storage mutations flow through
// the compat holder view, so tensor.data_ptr() stays aligned with storage()
// while preserving tensor-specific offsets for views.
void* data_ptr() const {
if (!tensor_.defined()) {
return nullptr;
}
return const_cast<void*>(tensor_.data());
}
template <typename T>
T* data_ptr() const {
return static_cast<T*>(data_ptr());
}
const void* const_data_ptr() const { return data_ptr(); }
template <typename T>
const T* const_data_ptr() const;
void* mutable_data_ptr() const { return data_ptr(); }
template <typename T>
T* mutable_data_ptr() const;
int64_t stride(int64_t dim) const {
if (dim < 0) {
dim += tensor_.strides().size();
}
return tensor_.strides()[static_cast<int>(dim)];
}
c10::SymInt sym_stride(int64_t dim) const {
return static_cast<c10::SymInt>(stride(dim));
}
c10::IntArrayRef strides() const {
return compat::_PD_PhiDDimToIntArrayRef(tensor_.strides());
}
c10::SymIntArrayRef sym_strides() const {
return c10::SymIntArrayRef(
reinterpret_cast<const c10::SymInt*>(strides().data()),
strides().size());
}
int64_t size(int64_t dim) const {
if (dim < 0) {
dim += tensor_.dims().size();
}
return tensor_.dims()[static_cast<int>(dim)];
}
c10::SymInt sym_size(int64_t dim) const {
return static_cast<c10::SymInt>(size(dim));
}
c10::IntArrayRef sizes() const {
return compat::_PD_PhiDDimToIntArrayRef(tensor_.dims());
}
c10::SymIntArrayRef sym_sizes() const {
return c10::SymIntArrayRef(
reinterpret_cast<const c10::SymInt*>(sizes().data()), sizes().size());
}
int64_t numel() const { return tensor_.numel(); }
c10::SymInt sym_numel() const { return static_cast<c10::SymInt>(numel()); }
caffe2::TypeMeta dtype() const {
return caffe2::TypeMeta::fromScalarType(
compat::_PD_PhiDataTypeToAtenScalarType(tensor_.dtype()));
}
c10::Device device() const { return c10::Device(tensor_.place()); }
c10::DeviceIndex get_device() const {
return c10::Device(tensor_.place()).index();
}
int64_t dim() const { return tensor_.dims().size(); }
int64_t ndimension() const { return dim(); }
at::TensorBase contiguous(
c10::MemoryFormat memory_format = c10::MemoryFormat::Contiguous) const {
PD_CHECK(memory_format == c10::MemoryFormat::Contiguous,
"`MemoryFormat` other than Contiguous");
return TensorBase(tensor_.contiguous());
}
bool is_contiguous(
at::MemoryFormat memory_format = at::MemoryFormat::Contiguous) const {
PD_CHECK(memory_format == c10::MemoryFormat::Contiguous,
"`MemoryFormat` other than Contiguous");
return tensor_.is_contiguous();
}
bool is_non_overlapping_and_dense() const {
if (numel() <= 1) {
return true;
}
if (tensor_.is_contiguous()) {
return true;
}
// For non-contiguous tensors, verify sorted strides form a valid dense
// layout without gaps or overlaps.
auto sizes_vec = sizes();
auto strides_vec = strides();
int64_t ndim = dim();
std::vector<int64_t> perm(ndim);
for (int64_t i = 0; i < ndim; ++i) {
perm[i] = i;
}
std::sort(perm.begin(), perm.end(), [&](int64_t a, int64_t b) {
return strides_vec[a] < strides_vec[b];
});
int64_t expected_stride = 1;
for (int64_t i = 0; i < ndim; ++i) {
int64_t dim_idx = perm[i];
if (sizes_vec[dim_idx] == 0) {
return true;
}
if (sizes_vec[dim_idx] == 1) {
continue;
}
if (strides_vec[dim_idx] != expected_stride) {
return false;
}
expected_stride *= sizes_vec[dim_idx];
}
return true;
}
bool is_contiguous_or_false(
at::MemoryFormat memory_format = at::MemoryFormat::Contiguous) const {
PD_CHECK(memory_format == c10::MemoryFormat::Contiguous,
"`MemoryFormat` other than Contiguous");
return tensor_.is_contiguous();
}
c10::ScalarType scalar_type() const {
return compat::_PD_PhiDataTypeToAtenScalarType(tensor_.dtype());
}
bool has_names() const {
// In PyTorch, has_names() is used to check if any dimension has names.
// In Paddle, we don't support named dimension yet, so always return false.
return false;
}
TensorOptions options() const {
return TensorOptions().dtype(dtype()).device(device()).layout(layout());
}
const TensorBase& fill_(const at::Scalar& scalar) const {
paddle::experimental::fill_(const_cast<PaddleTensor&>(tensor_), scalar);
return *this;
}
const TensorBase& zero_() const {
paddle::experimental::fill_(const_cast<PaddleTensor&>(tensor_), 0.0);
return *this;
}
at::TensorBase to(
at::TensorOptions options = {},
bool non_blocking = false,
bool copy = false,
std::optional<at::MemoryFormat> memory_format = std::nullopt) const {
if (options.device_opt().has_value()) {
PADDLE_THROW(common::errors::Unimplemented(
"The `to` method with device option is not supported yet."));
}
if (memory_format.has_value()) {
PADDLE_THROW(common::errors::Unimplemented(
"The `to` method with memory_format option is not supported yet."));
}
return TensorBase(paddle::experimental::cast(
tensor_, compat::_PD_AtenScalarTypeToPhiDataType(options.dtype())));
}
bool is_complex() const { return at::isComplexType(this->scalar_type()); }
bool is_floating_point() const {
return at::isFloatingType(this->scalar_type());
}
bool is_signed() const { return at::isSignedType(this->scalar_type()); }
bool is_cpu() const { return phi::is_cpu_place(tensor_.place()); }
bool is_cuda() const { return phi::is_gpu_place(tensor_.place()); }
bool is_sparse() const {
return tensor_.is_sparse_coo_tensor() || tensor_.is_sparse_csr_tensor();
}
bool is_sparse_csr() const { return tensor_.is_sparse_csr_tensor(); }
inline size_t nbytes() const {
PD_CHECK(
((tensor_.layout() != common::DataLayout::SPARSE_COO) &&
(tensor_.layout() != common::DataLayout::SPARSE_CSR)),
"nbytes is not defined for sparse tensors. If you want the size of "
"the constituent "
"tensors, add the nbytes of the indices and values. If you want the "
"size of the "
"equivalent dense tensor, multiply numel() by element_size()");
return tensor_.numel() * SizeOf(tensor_.dtype());
}
size_t itemsize() const { return SizeOf(tensor_.dtype()); }
int64_t element_size() const {
return static_cast<int64_t>(SizeOf(tensor_.dtype()));
}
bool defined() const { return tensor_.defined(); }
Layout layout() const {
switch (tensor_.layout()) {
case common::DataLayout::STRIDED:
case common::DataLayout::NCHW:
case common::DataLayout::NHWC:
case common::DataLayout::NCDHW:
case common::DataLayout::NDHWC:
return c10::kStrided;
case common::DataLayout::SPARSE_COO:
return c10::kSparse;
case common::DataLayout::SPARSE_CSR:
return c10::kSparseCsr;
case common::DataLayout::ONEDNN:
return c10::kMkldnn;
default:
return c10::kStrided;
}
}
void reset() {
tensor_.reset();
storage_.reset();
}
int64_t storage_offset() const {
// Paddle DenseTensor stores offset in meta_.offset (in bytes)
// We need to convert to element offset
auto dense_tensor =
std::dynamic_pointer_cast<phi::DenseTensor>(tensor_.impl());
if (dense_tensor) {
size_t byte_offset = dense_tensor->meta().offset;
size_t element_size = SizeOf(tensor_.dtype());
return element_size > 0 ? static_cast<int64_t>(byte_offset / element_size)
: 0;
}
return 0;
}
c10::SymInt sym_storage_offset() const {
return c10::SymInt(storage_offset());
}
bool has_storage() const {
SyncStorageFromTensor();
return tensor_.defined() && storage_ && storage_->valid();
}
// Returns a Storage handle backed by the shared StorageImpl for this tensor.
const c10::Storage& storage() const {
SyncStorageFromTensor();
static const c10::Storage kEmptyStorage;
return storage_ ? *storage_ : kEmptyStorage;
}
bool is_alias_of(const at::TensorBase& other) const {
return this->storage().is_alias_of(other.storage());
}
private:
template <typename DenseT>
static auto MaybeResetHolder(DenseT* dense,
const std::shared_ptr<phi::Allocation>& holder,
int)
-> decltype(dense->ResetHolder(holder), void()) {
dense->ResetHolder(holder);
}
static void MaybeResetHolder(phi::DenseTensor* dense,
const std::shared_ptr<phi::Allocation>& holder,
long) { // NOLINT
TORCH_CHECK(dense != nullptr, "DenseTensor must not be null");
// External custom-kernel builds do not expose ResetHolder(), but Holder()
// still returns the live holder reference used by DenseTensor.
if (dense->numel() == 0) {
auto& meta = const_cast<phi::DenseTensorMeta&>(dense->meta());
meta.offset = 0;
const_cast<std::shared_ptr<phi::Allocation>&>(dense->Holder()) = holder;
return;
}
if (dense->Holder() && dense->meta().is_contiguous()) {
TORCH_CHECK(holder != nullptr, "Holder must not be null.");
const auto required_bytes =
dense->numel() * static_cast<int64_t>(phi::SizeOf(dense->dtype())) +
static_cast<int64_t>(dense->meta().offset);
TORCH_CHECK(required_bytes <= static_cast<int64_t>(holder->size()),
"The size of Holder is not enough to store the Tensor.");
}
const_cast<std::shared_ptr<phi::Allocation>&>(dense->Holder()) = holder;
}
void InitStorage() { SyncStorageFromTensor(); }
static std::shared_ptr<c10::Storage> GetOrCreateCanonicalStorage(
c10::Storage&& live_storage) {
auto impl = live_storage.get_impl();
if (!impl) {
return std::make_shared<c10::Storage>(std::move(live_storage));
}
static std::mutex registry_mu;
static std::unordered_map<c10::StorageImpl*, std::weak_ptr<c10::Storage>>
registry;
std::lock_guard<std::mutex> guard(registry_mu);
auto it = registry.find(impl.get());
if (it != registry.end()) {
if (auto cached = it->second.lock()) {
return cached;
}
registry.erase(it);
}
auto created = std::make_shared<c10::Storage>(std::move(live_storage));
registry.emplace(impl.get(), created);
return created;
}
void SyncStorageFromTensor() const {
auto dense = std::dynamic_pointer_cast<phi::DenseTensor>(tensor_.impl());
if (!dense) {
storage_.reset();
return;
}
auto holder = dense->Holder();
if (!holder) {
storage_.reset();
return;
}
c10::Storage live_storage = c10::Storage::createTensorStorage(holder);
auto compat_holder = live_storage.ensureTensorHolder();
if (holder != compat_holder) {
MaybeResetHolder(dense.get(), compat_holder, 0);
}
if (!storage_ || storage_->get_impl() != live_storage.get_impl()) {
storage_ = GetOrCreateCanonicalStorage(std::move(live_storage));
}
}
public:
// Return a `TensorAccessor` for CPU `Tensor`s. You have to specify scalar
// type and
// dimension.
template <typename T, size_t N>
TensorAccessor<T, N> accessor() const& {
static_assert(
N > 0,
"accessor is used for indexing tensor, for scalars use *data_ptr<T>()");
TORCH_CHECK(dim() == N,
"TensorAccessor expected ",
N,
" dims but tensor has ",
dim());
T* ptr = nullptr;
if constexpr (std::is_const_v<T>) {
ptr = const_data_ptr<T>();
} else {
ptr = mutable_data_ptr<T>();
}
return TensorAccessor<T, N>(ptr, sizes().data(), strides().data());
}
template <typename T, size_t N>
TensorAccessor<T, N> accessor() && = delete;
// Return a `GenericPackedTensorAccessor` for CUDA `Tensor`s. You have to
// specify scalar type and dimension. You can optionally specify
// RestrictPtrTraits as a template parameter to cast the data pointer to a
// __restrict__ pointer. In order to use this, your CUDA kernel has to take a
// corresponding GenericPackedTensorAccessor as an argument.
template <typename T,
size_t N,
template <typename U> class PtrTraits = DefaultPtrTraits,
typename index_t = int64_t>
GenericPackedTensorAccessor<T, N, PtrTraits, index_t>
generic_packed_accessor() const& {
static_assert(
N > 0,
"accessor is used for indexing tensor, for scalars use *data_ptr<T>()");
TORCH_CHECK(dim() == N,
"TensorAccessor expected ",
N,
" dims but tensor has ",
dim());
T* ptr = nullptr;
if constexpr (std::is_const_v<T>) {
ptr = const_data_ptr<T>();
} else {
ptr = mutable_data_ptr<T>();
}
return GenericPackedTensorAccessor<T, N, PtrTraits, index_t>(
static_cast<typename PtrTraits<T>::PtrType>(ptr),
sizes().data(),
strides().data());
}
template <typename T,
size_t N,
template <typename U> class PtrTraits = DefaultPtrTraits,
typename index_t = int64_t>
GenericPackedTensorAccessor<T, N> generic_packed_accessor() && = delete;
template <typename T,
size_t N,
template <typename U> class PtrTraits = DefaultPtrTraits>
PackedTensorAccessor32<T, N, PtrTraits> packed_accessor32() const& {
TORCH_CHECK(
numel() <= static_cast<int64_t>(std::numeric_limits<int32_t>::max()),
"numel needs to be smaller than int32_t max; otherwise, please use "
"packed_accessor64");
return generic_packed_accessor<T, N, PtrTraits, int32_t>();
}
template <typename T,
size_t N,
template <typename U> class PtrTraits = DefaultPtrTraits>
PackedTensorAccessor32<T, N, PtrTraits> packed_accessor32() && = delete;
template <typename T,
size_t N,
template <typename U> class PtrTraits = DefaultPtrTraits>
PackedTensorAccessor64<T, N, PtrTraits> packed_accessor64() const& {
return generic_packed_accessor<T, N, PtrTraits, int64_t>();
}
template <typename T,
size_t N,
template <typename U> class PtrTraits = DefaultPtrTraits>
PackedTensorAccessor64<T, N, PtrTraits> packed_accessor64() && = delete;
const PaddleTensor& _PD_GetInner() const& { return tensor_; }
PaddleTensor& _PD_GetInner() & { return tensor_; }
PaddleTensor&& _PD_GetInner() && { return std::move(tensor_); }
protected:
PaddleTensor tensor_;
// Cache a canonical Storage object shared by wrappers that reference the
// same StorageImpl. This prevents independently-constructed wrappers around
// one tensor impl from inflating Storage::use_count().
mutable std::shared_ptr<c10::Storage> storage_;
};
} // namespace at
@@ -0,0 +1,846 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/TensorIndexing.h>
#include <ATen/core/TensorBase.h>
#include <c10/core/Backend.h>
#include <c10/core/List.h>
#include <c10/core/Scalar.h>
#include <c10/core/ScalarType.h>
#include <c10/core/Stream.h>
#include <c10/core/SymIntArrayRef.h>
#include <c10/util/OptionalArrayRef.h>
#include "paddle/phi/api/include/api.h"
#include "paddle/phi/api/include/tensor.h"
#include "paddle/phi/common/int_array.h"
#include "paddle/phi/common/scalar.h"
#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/core/memory/malloc.h"
#ifdef PADDLE_WITH_HIP
#include <hip/hip_runtime.h>
#elif defined(PADDLE_WITH_CUDA)
#include <cuda_runtime_api.h>
#endif
#include <limits>
#include <optional>
#include <type_traits>
#include <utility>
#include <vector>
#include "paddle/common/ddim.h"
#include "paddle/phi/common/place.h"
#include "paddle/phi/core/enforce.h"
namespace at {
class Tensor;
// Type aliases for ATen compatibility
using Scalar = c10::Scalar;
using TensorOptions = c10::TensorOptions;
using MemoryFormat = c10::MemoryFormat;
using IntArrayRef = c10::IntArrayRef;
using OptionalIntArrayRef = c10::OptionalIntArrayRef;
using ScalarType = c10::ScalarType;
using TensorList = c10::ArrayRef<Tensor>;
using ITensorListRef = c10::ArrayRef<Tensor>;
} // namespace at
namespace at { // NOLINT(build/namespaces)
using PaddleTensor = paddle::Tensor;
using PaddlePlace = phi::Place;
// Stub for DimnameList (not supported in Paddle)
using DimnameList = c10::ArrayRef<std::string>;
using Stream = c10::Stream;
class Tensor : public TensorBase {
public:
Tensor() = default;
Tensor(const PaddleTensor& tensor) : TensorBase(tensor){}; // NOLINT
Tensor(const Tensor& tensor) = default;
Tensor(Tensor&& tensor) = default;
// Implicitly move-constructible from TensorBase, but must be explicit to
// increase refcount
explicit Tensor(const TensorBase& base) : TensorBase(base) {} // NOLINT
/*implicit*/ Tensor(TensorBase&& base) // NOLINT
: TensorBase(std::move(base)) {}
Tensor& operator=(const PaddleTensor& x) & noexcept {
tensor_ = x;
return *this;
}
Tensor& operator=(const TensorBase& x) & noexcept {
const PaddleTensor& inner = x._PD_GetInner();
tensor_ = inner;
return *this;
}
Tensor& operator=(PaddleTensor&& x) & noexcept {
tensor_ = std::move(x);
return *this;
}
Tensor& operator=(TensorBase&& x) & noexcept {
tensor_ = std::move(x)._PD_GetInner();
return *this;
}
Tensor& operator=(const Tensor& x) & noexcept {
return operator=(static_cast<const TensorBase&>(x));
}
Tensor& operator=(Tensor&& x) & noexcept {
return operator=(static_cast<TensorBase&&>(x));
}
Tensor& operator=(const Scalar& v) && {
fill_(v);
return *this;
}
Tensor& operator=(const Tensor& rhs) && {
copy_(rhs);
return *this;
}
Tensor& operator=(Tensor&& rhs) && {
copy_(rhs);
return *this;
}
template <typename T>
T* data() const {
return data_ptr<T>();
}
Tensor toBackend(c10::Backend b) const {
switch (b) {
case c10::Backend::CPU:
return tensor_.copy_to(PaddlePlace(phi::AllocationType::CPU), true);
case c10::Backend::CUDA:
return tensor_.copy_to(paddle::DefaultGPUPlace(), true);
case c10::Backend::XPU:
return tensor_.copy_to(paddle::DefaultXPUPlace(), true);
case c10::Backend::IPU:
return tensor_.copy_to(PaddlePlace(phi::IPUPlace()), true);
default:
PD_CHECK(false, "Unsupported backend");
}
return tensor_;
}
Tensor cpu() const {
PaddlePlace place(phi::AllocationType::CPU);
return tensor_.copy_to(place, true);
}
Tensor cuda() const {
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
auto place = paddle::DefaultGPUPlace();
return tensor_.copy_to(place, true);
#elif defined(PADDLE_WITH_XPU)
return tensor_.copy_to(paddle::DefaultXPUPlace(), true);
#elif defined(PADDLE_WITH_CUSTOM_DEVICE)
return tensor_.copy_to(paddle::DefaultCustomPlace(), true);
#else
PD_THROW(
"cuda() is not supported: no GPU/XPU/Custom device enabled "
"in this build.");
#endif
}
using TensorBase::stride;
using TensorBase::size;
at::Tensor to(
at::TensorOptions options = {},
bool non_blocking = false,
bool copy = false,
::std::optional<at::MemoryFormat> memory_format = ::std::nullopt) const;
at::Tensor to(::std::optional<at::ScalarType> dtype,
::std::optional<at::Layout> layout,
::std::optional<at::Device> device,
::std::optional<bool> pin_memory,
bool non_blocking,
bool copy,
::std::optional<at::MemoryFormat> memory_format) const;
at::Tensor to(
at::Device device,
at::ScalarType dtype,
bool non_blocking = false,
bool copy = false,
::std::optional<at::MemoryFormat> memory_format = ::std::nullopt) const;
at::Tensor to(
at::ScalarType dtype,
bool non_blocking = false,
bool copy = false,
::std::optional<at::MemoryFormat> memory_format = ::std::nullopt) const;
at::Tensor to(
const at::Tensor& other,
bool non_blocking = false,
bool copy = false,
::std::optional<at::MemoryFormat> memory_format = ::std::nullopt) const;
Tensor meta() const {
PD_THROW("`meta()` is not supported in this Paddle build.");
}
at::Scalar item() const;
template <typename T>
T item() const;
bool equal(const at::Tensor& other) const;
// Clamp functions
at::Tensor clamp(
const ::std::optional<at::Scalar>& min,
const ::std::optional<at::Scalar>& max = ::std::nullopt) const;
at::Tensor clamp(const ::std::optional<at::Tensor>& min = {},
const ::std::optional<at::Tensor>& max = {}) const;
at::Tensor& clamp_(
const ::std::optional<at::Scalar>& min,
const ::std::optional<at::Scalar>& max = ::std::nullopt) const;
at::Tensor& clamp_(const ::std::optional<at::Tensor>& min = {},
const ::std::optional<at::Tensor>& max = {}) const;
at::Tensor clamp_max(const at::Scalar& max) const;
at::Tensor clamp_max(const at::Tensor& max) const;
at::Tensor& clamp_max_(const at::Scalar& max) const;
at::Tensor& clamp_max_(const at::Tensor& max) const;
at::Tensor clamp_min(const at::Scalar& min) const;
at::Tensor clamp_min(const at::Tensor& min) const;
at::Tensor& clamp_min_(const at::Scalar& min) const;
at::Tensor& clamp_min_(const at::Tensor& min) const;
// as_strided: Create a tensor view with custom size, stride, and
// storage_offset
at::Tensor as_strided(
at::IntArrayRef size,
at::IntArrayRef stride,
::std::optional<int64_t> storage_offset = ::std::nullopt) const;
// as_strided_: Inplace version
const at::Tensor& as_strided_(
at::IntArrayRef size,
at::IntArrayRef stride,
::std::optional<int64_t> storage_offset = ::std::nullopt) const;
// as_strided_scatter: Scatter src into a strided view
at::Tensor as_strided_scatter(
const at::Tensor& src,
at::IntArrayRef size,
at::IntArrayRef stride,
::std::optional<int64_t> storage_offset = ::std::nullopt) const;
// Standard deviation functions
Tensor std(bool unbiased) const;
Tensor std(at::OptionalIntArrayRef dim,
bool unbiased,
bool keepdim = false) const;
Tensor std(at::OptionalIntArrayRef dim = ::std::nullopt,
const ::std::optional<at::Scalar>& correction = ::std::nullopt,
bool keepdim = false) const;
Tensor std(int dim) const { return std(at::IntArrayRef{dim}); }
Tensor tensor_data() const {
PaddleTensor result;
if (tensor_.initialized()) {
auto src_impl = tensor_.impl();
auto* src_tensor =
std::dynamic_pointer_cast<phi::DenseTensor>(src_impl).get();
if (src_tensor && src_tensor->meta().is_contiguous()) {
result.set_impl(std::make_shared<phi::DenseTensor>());
auto* dst_tensor =
std::dynamic_pointer_cast<phi::DenseTensor>(result.impl()).get();
dst_tensor->ShareDataWith(*src_tensor);
} else {
result = paddle::experimental::assign(tensor_);
}
}
// For uninitialized tensor, return an uninitialized tensor (no assign
// needed)
return Tensor(result);
}
Tensor variable_data() const {
PaddleTensor result;
if (tensor_.initialized()) {
auto src_impl = tensor_.impl();
auto* src_tensor =
std::dynamic_pointer_cast<phi::DenseTensor>(src_impl).get();
if (src_tensor && src_tensor->meta().is_contiguous()) {
result.set_impl(std::make_shared<phi::DenseTensor>());
auto* dst_tensor =
std::dynamic_pointer_cast<phi::DenseTensor>(result.impl()).get();
dst_tensor->ShareDataWith(*src_tensor);
} else {
result = paddle::experimental::assign(tensor_);
}
}
// For uninitialized tensor, return an uninitialized tensor (no assign
// needed)
return Tensor(result);
}
// index: Get values at specified tensor indices
at::Tensor index(const c10::List<::std::optional<at::Tensor>>& indices) const;
// index_put_: Set values at specified indices in-place
at::Tensor& index_put_(const c10::List<::std::optional<at::Tensor>>& indices,
const at::Tensor& values,
bool accumulate = false) const;
// index_put: Non-inplace version of index_put_
at::Tensor index_put(const c10::List<::std::optional<at::Tensor>>& indices,
const at::Tensor& values,
bool accumulate = false) const;
Tensor toType(ScalarType t) const {
return Tensor(paddle::experimental::cast(
tensor_, compat::_PD_AtenScalarTypeToPhiDataType(t)));
}
at::Tensor contiguous(
c10::MemoryFormat memory_format = c10::MemoryFormat::Contiguous) const {
PD_CHECK(memory_format == c10::MemoryFormat::Contiguous,
"`MemoryFormat` other than Contiguous");
return tensor_.contiguous();
}
at::Tensor flatten(int64_t start_dim = 0, int64_t end_dim = -1) const;
at::Tensor unflatten(int64_t dim, at::IntArrayRef sizes) const;
at::Tensor unflatten_symint(int64_t dim, c10::SymIntArrayRef sizes) const;
Tensor& fill_(const at::Scalar& value) const {
paddle::experimental::fill_(const_cast<PaddleTensor&>(tensor_), value);
return const_cast<at::Tensor&>(*this);
}
Tensor& zero_() const {
paddle::experimental::fill_(const_cast<PaddleTensor&>(tensor_), 0.0);
return const_cast<at::Tensor&>(*this);
}
bool is_pinned(::std::optional<c10::Device> device = ::std::nullopt) const {
if (device.has_value()) {
phi::enforce::ThrowWarnInternal(
"The argument 'device' of Tensor.is_pinned() is deprecated. "
"Please do not pass this argument.");
}
const PaddlePlace place = tensor_.place();
const bool is_gpu_pinned = phi::is_cuda_pinned_place(place);
const bool is_xpu_pinned = phi::is_xpu_pinned_place(place);
// Keep parity with PyTorch behavior: only host tensors are pinnable.
if (!(phi::is_cpu_place(place) || is_gpu_pinned || is_xpu_pinned)) {
return false;
}
if (!device.has_value()) {
return is_gpu_pinned || is_xpu_pinned;
}
const auto device_type = device.value().type();
if (device_type == c10::DeviceType::CUDA) {
return is_gpu_pinned;
}
if (device_type == c10::DeviceType::XPU) {
return is_xpu_pinned;
}
// CPU and non-accelerator devices are not valid pinned backends.
return false;
}
Tensor pin_memory(
::std::optional<c10::Device> device = ::std::nullopt) const {
if (device.has_value()) {
phi::enforce::ThrowWarnInternal(
"The argument 'device' of Tensor.pin_memory() is deprecated. "
"Please do not pass this argument.");
}
if (is_pinned(device)) {
return *this;
}
const PaddlePlace current_place = tensor_.place();
if (!phi::is_cpu_place(current_place)) {
PD_THROW("cannot pin '" + this->toString() +
"', only dense CPU tensors can be pinned");
}
PaddlePlace pinned_place;
if (device.has_value()) {
const auto device_type = device.value().type();
if (device_type == c10::DeviceType::CUDA) {
pinned_place = phi::Place(phi::GPUPinnedPlace());
} else if (device_type == c10::DeviceType::XPU) {
pinned_place = phi::Place(phi::XPUPinnedPlace());
} else {
PD_THROW("pin_memory device type must be an accelerator (GPU/XPU)");
}
} else {
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
pinned_place = phi::Place(phi::GPUPinnedPlace());
#elif defined(PADDLE_WITH_XPU)
pinned_place = phi::Place(phi::XPUPinnedPlace());
#else
PD_THROW("pin_memory is not supported: no GPU/XPU backend enabled");
#endif
}
return tensor_.copy_to(pinned_place, true);
}
at::Tensor narrow_copy(int64_t dim, int64_t start, int64_t length) const;
at::Tensor narrow_copy_symint(int64_t dim,
c10::SymInt start,
c10::SymInt length) const;
at::Tensor narrow(int64_t dim, int64_t start, int64_t length) const;
at::Tensor narrow_symint(int64_t dim,
c10::SymInt start,
c10::SymInt length) const;
at::Tensor narrow(int64_t dim, const at::Tensor& start, int64_t length) const;
at::Tensor narrow_symint(int64_t dim,
const at::Tensor& start,
c10::SymInt length) const;
at::Tensor reshape(at::IntArrayRef shape) const;
at::Tensor transpose(int64_t dim0, int64_t dim1) const;
at::Tensor& transpose_(int64_t dim0, int64_t dim1) const;
at::Tensor permute(at::IntArrayRef dims) const;
at::Tensor reciprocal() const;
at::Tensor& reciprocal_() const;
at::Tensor detach() const;
at::Tensor& detach_() const;
at::Tensor select(int64_t dim, int64_t index) const;
at::Tensor select_symint(int64_t dim, c10::SymInt index) const;
at::Tensor& copy_(const at::Tensor& src, bool non_blocking = false) const {
const_cast<PaddleTensor&>(tensor_).copy_(
src._PD_GetInner(), tensor_.place(), /*blocking=*/!non_blocking);
return const_cast<at::Tensor&>(*this);
}
at::Tensor view(at::IntArrayRef size) const;
at::Tensor view(at::ScalarType dtype) const;
at::Tensor squeeze() const;
at::Tensor squeeze(int64_t dim) const;
at::Tensor squeeze(at::IntArrayRef dim) const;
at::Tensor& squeeze_() const;
at::Tensor& squeeze_(int64_t dim) const;
at::Tensor& squeeze_(at::IntArrayRef dim) const;
at::Tensor unsqueeze(int64_t dim) const;
at::Tensor& unsqueeze_(int64_t dim) const;
at::Tensor sum(::std::optional<at::ScalarType> dtype = ::std::nullopt) const;
at::Tensor sum(at::OptionalIntArrayRef dim,
bool keepdim = false,
::std::optional<at::ScalarType> dtype = ::std::nullopt) const;
at::Tensor t() const;
at::Tensor& t_() const;
at::Tensor view_as(const at::Tensor& other) const;
at::Tensor coalesce() const;
bool is_coalesced() const;
int64_t _nnz() const;
at::Tensor _values() const;
bool is_variable() const noexcept { return true; }
at::Tensor index_select(int64_t dim, const at::Tensor& index) const {
return Tensor(
paddle::experimental::index_select(tensor_, index._PD_GetInner(), dim));
}
at::Tensor masked_select(const at::Tensor& mask) const;
std::vector<at::Tensor> tensor_split(int64_t sections, int64_t dim) const;
std::vector<at::Tensor> tensor_split_symint(c10::SymInt sections,
int64_t dim) const;
std::vector<at::Tensor> tensor_split(at::IntArrayRef indices,
int64_t dim) const;
std::vector<at::Tensor> tensor_split_symint(c10::SymIntArrayRef indices,
int64_t dim) const;
std::vector<at::Tensor> tensor_split(
const at::Tensor& tensor_indices_or_sections, int64_t dim) const;
std::vector<at::Tensor> split(int64_t split_size, int64_t dim) const;
std::vector<at::Tensor> split_symint(c10::SymInt split_size,
int64_t dim) const;
std::vector<at::Tensor> split(at::IntArrayRef split_sizes, int64_t dim) const;
std::vector<at::Tensor> split_symint(c10::SymIntArrayRef split_sizes,
int64_t dim) const;
std::vector<at::Tensor> unsafe_split(int64_t split_size,
int64_t dim = 0) const;
std::vector<at::Tensor> unsafe_split_symint(c10::SymInt split_size,
int64_t dim = 0) const;
std::vector<at::Tensor> split_with_sizes(at::IntArrayRef split_sizes,
int64_t dim = 0) const;
std::vector<at::Tensor> split_with_sizes_symint(
c10::SymIntArrayRef split_sizes, int64_t dim = 0) const;
std::vector<at::Tensor> unsafe_split_with_sizes(at::IntArrayRef split_sizes,
int64_t dim = 0) const;
std::vector<at::Tensor> unsafe_split_with_sizes_symint(
c10::SymIntArrayRef split_sizes, int64_t dim = 0) const;
std::vector<at::Tensor> hsplit(int64_t sections) const;
std::vector<at::Tensor> hsplit(at::IntArrayRef indices) const;
std::vector<at::Tensor> vsplit(int64_t sections) const;
std::vector<at::Tensor> vsplit(at::IntArrayRef indices) const;
std::vector<at::Tensor> dsplit(int64_t sections) const;
std::vector<at::Tensor> dsplit(at::IntArrayRef indices) const;
at::Tensor bitwise_right_shift(const Scalar& other) const {
return Tensor(paddle::experimental::bitwise_right_shift(
tensor_, paddle::experimental::full({}, other, other.dtype())));
}
at::Tensor slice(int64_t dim = 0,
::std::optional<int64_t> start = ::std::nullopt,
::std::optional<int64_t> end = ::std::nullopt,
int64_t step = 1) const;
at::Tensor index(ArrayRef<at::indexing::TensorIndex> indices) const;
inline at::Tensor index(
std::initializer_list<at::indexing::TensorIndex> indices) const {
return index(ArrayRef<at::indexing::TensorIndex>(indices));
}
Tensor& index_put_(ArrayRef<at::indexing::TensorIndex> indices,
Tensor const& rhs);
Tensor& index_put_(ArrayRef<at::indexing::TensorIndex> indices,
const Scalar& v);
Tensor& index_put_(std::initializer_list<at::indexing::TensorIndex> indices,
Tensor const& rhs);
Tensor& index_put_(std::initializer_list<at::indexing::TensorIndex> indices,
const Scalar& v);
at::Tensor& floor_divide_(const at::Scalar& other) const {
paddle::experimental::floor_divide_(
const_cast<PaddleTensor&>(tensor_),
paddle::experimental::full({}, other, other.dtype()));
return const_cast<at::Tensor&>(*this);
}
// Paddle Tensor has no storage_offset, so we add it here, and it is always
// 0.
// int64_t storage_offset() const { return storage_offset_; }
inline Tensor clone(
::std::optional<at::MemoryFormat> memory_format = ::std::nullopt) const {
(void)memory_format;
PaddleTensor cloned_tensor = paddle::experimental::assign(tensor_);
return Tensor(cloned_tensor);
}
// all: Check if all elements are true (non-zero)
at::Tensor all() const;
at::Tensor all(int64_t dim, bool keepdim = false) const;
at::Tensor all(at::OptionalIntArrayRef dim, bool keepdim = false) const;
// allclose: Check if two tensors are close to each other
bool allclose(const at::Tensor& other,
double rtol = 1e-05,
double atol = 1e-08,
bool equal_nan = false) const;
at::Tensor abs() const;
at::Tensor& abs_() const;
at::Tensor absolute() const { return abs(); }
at::Tensor& absolute_() const { return abs_(); }
Tensor operator[](int64_t index) const {
// Use as_strided to create a view (shares storage with original tensor)
// This allows fill_ to modify the original tensor
int64_t numel = tensor_.numel();
if (numel == 0) {
PD_THROW("operator[]: cannot index empty tensor");
}
// Handle negative index
if (index < 0) {
index += tensor_.dims()[0];
}
// Check bounds
if (index < 0 || index >= tensor_.dims()[0]) {
PD_THROW("operator[]: index ",
index,
" out of range for tensor of size ",
tensor_.dims(),
" at dimension 0");
}
// For 1D tensor: create a scalar view (0-dim tensor) with proper offset
// For multi-D tensor: create a view of the row at index
std::vector<int64_t> new_sizes;
std::vector<int64_t> new_strides;
auto dims = tensor_.dims();
auto stride = tensor_.strides();
// Skip the first dimension (dim 0)
for (int i = 1; i < dims.size(); ++i) {
new_sizes.push_back(dims[i]);
new_strides.push_back(stride[i]);
}
// Calculate storage offset
int64_t storage_offset = index * stride[0];
return as_strided(c10::IntArrayRef(new_sizes),
c10::IntArrayRef(new_strides),
storage_offset);
}
void record_stream(at::Stream s) const;
Tensor var(int dim) const { return var(at::IntArrayRef{dim}, true, false); }
Tensor var(bool unbiased) const {
std::vector<int64_t> empty_dims;
double correction = unbiased ? 1.0 : 0.0;
return var_impl(empty_dims, correction, false);
}
Tensor var(at::OptionalIntArrayRef dim,
bool unbiased,
bool keepdim = false) const {
// Convert unbiased to correction: unbiased=True means correction=1
double correction = unbiased ? 1.0 : 0.0;
std::vector<int64_t> dims_vec;
if (dim.has_value() && dim.value().size() > 0) {
dims_vec.assign(dim.value().begin(), dim.value().end());
}
return var_impl(dims_vec, correction, keepdim);
}
Tensor var(at::OptionalIntArrayRef dim = ::std::nullopt,
const ::std::optional<at::Scalar>& correction = ::std::nullopt,
bool keepdim = false) const {
double correction_value = 1.0;
if (correction.has_value()) {
const at::Scalar& scalar = correction.value();
correction_value = scalar.to<double>();
}
std::vector<int64_t> dims_vec;
if (dim.has_value() && dim.value().size() > 0) {
dims_vec.assign(dim.value().begin(), dim.value().end());
}
return var_impl(dims_vec, correction_value, keepdim);
}
private:
Tensor var_impl(const std::vector<int64_t>& dims_vec,
double correction_value,
bool keepdim) const {
phi::IntArray dims_int_array(dims_vec);
PaddleTensor mean_tensor;
if (dims_vec.empty()) {
mean_tensor = paddle::experimental::mean(
tensor_, phi::IntArray(std::vector<int64_t>{}), true);
} else {
mean_tensor = paddle::experimental::mean(tensor_, dims_int_array, true);
}
PaddleTensor diff = paddle::experimental::subtract(tensor_, mean_tensor);
PaddleTensor diff_squared = paddle::experimental::multiply(diff, diff);
PaddleTensor sum_squared_diff;
if (dims_vec.empty()) {
sum_squared_diff =
paddle::experimental::sum(diff_squared,
phi::IntArray(std::vector<int64_t>{}),
diff_squared.dtype(),
keepdim);
} else {
sum_squared_diff = paddle::experimental::sum(
diff_squared, dims_int_array, diff_squared.dtype(), keepdim);
}
int64_t n = tensor_.numel();
if (!dims_vec.empty()) {
n = 1;
for (int64_t d : dims_vec) {
int64_t dim_idx = d < 0 ? d + tensor_.dims().size() : d;
if (dim_idx >= 0 &&
dim_idx < static_cast<int64_t>(tensor_.dims().size())) {
n *= tensor_.dims()[dim_idx];
}
}
}
double corrected_n = static_cast<double>(n) - correction_value;
if (corrected_n <= 0.0) {
corrected_n = static_cast<double>(n);
}
std::vector<int64_t> result_shape_vec;
for (int64_t i = 0; i < sum_squared_diff.dims().size(); ++i) {
result_shape_vec.push_back(sum_squared_diff.dims()[i]);
}
PaddleTensor correction_scalar =
paddle::experimental::full(phi::IntArray(result_shape_vec),
phi::Scalar(corrected_n),
sum_squared_diff.dtype(),
sum_squared_diff.place());
PaddleTensor result =
paddle::experimental::divide(sum_squared_diff, correction_scalar);
return Tensor(result);
}
public:
// Deprecated packed_accessor for compatibility with PyTorch
// Use packed_accessor32 or packed_accessor64 instead
template <typename T,
size_t N,
template <typename U> class PtrTraits = DefaultPtrTraits,
typename index_t = int64_t>
[[deprecated(
"packed_accessor is deprecated, use packed_accessor32 or "
"packed_accessor64 instead")]] GenericPackedTensorAccessor<T,
N,
PtrTraits,
index_t>
packed_accessor() const& {
return this->template generic_packed_accessor<T, N, PtrTraits, index_t>();
}
template <typename T,
size_t N,
template <typename U> class PtrTraits = DefaultPtrTraits,
typename index_t = int64_t>
[[deprecated(
"packed_accessor is deprecated, use packed_accessor32 or "
"packed_accessor64 instead")]] GenericPackedTensorAccessor<T,
N,
PtrTraits,
index_t>
packed_accessor() && = delete;
template <typename T>
using hook_return_void_t =
std::enable_if_t<std::is_void_v<std::invoke_result_t<T&, Tensor>>,
unsigned>;
template <typename T>
using hook_return_var_t =
std::enable_if_t<std::is_same_v<std::invoke_result_t<T&, Tensor>, Tensor>,
unsigned>;
// register_hook - throws exception for Paddle compatibility
// Paddle does not support gradient hooks
template <typename T>
hook_return_void_t<T> register_hook(T&&) const {
throw std::runtime_error(
"register_hook is not supported in Paddle, this is an ATen "
"compatibility API that is not available");
}
template <typename T>
hook_return_var_t<T> register_hook(T&&) const {
throw std::runtime_error(
"register_hook is not supported in Paddle, this is an ATen "
"compatibility API that is not available");
}
// any - returns true if any element is non-zero
Tensor any(int64_t dim, bool keepdim = false) const;
Tensor any(at::OptionalIntArrayRef dim, bool keepdim = false) const;
Tensor any() const;
// chunk - splits tensor into chunks
std::vector<Tensor> chunk(int64_t chunks, int64_t dim = 0) const;
// rename - stub for Paddle (Dimname not supported)
Tensor rename(::std::optional<at::DimnameList> names) const;
// new_empty - creates uninitialized tensor with same dtype/device
Tensor new_empty(at::IntArrayRef size, at::TensorOptions options = {}) const;
Tensor new_empty(at::IntArrayRef size,
::std::optional<at::ScalarType> dtype,
::std::optional<at::Layout> layout,
::std::optional<at::Device> device,
::std::optional<bool> pin_memory) const;
// new_full - creates tensor filled with fill_value
Tensor new_full(at::IntArrayRef size,
const at::Scalar& fill_value,
at::TensorOptions options = {}) const;
Tensor new_full(at::IntArrayRef size,
const at::Scalar& fill_value,
::std::optional<at::ScalarType> dtype,
::std::optional<at::Layout> layout,
::std::optional<at::Device> device,
::std::optional<bool> pin_memory) const;
// new_zeros - creates zero tensor
Tensor new_zeros(at::IntArrayRef size, at::TensorOptions options = {}) const;
Tensor new_zeros(at::IntArrayRef size,
::std::optional<at::ScalarType> dtype,
::std::optional<at::Layout> layout,
::std::optional<at::Device> device,
::std::optional<bool> pin_memory) const;
// new_ones - creates tensor filled with ones
Tensor new_ones(at::IntArrayRef size, at::TensorOptions options = {}) const;
Tensor new_ones(at::IntArrayRef size,
::std::optional<at::ScalarType> dtype,
::std::optional<at::Layout> layout,
::std::optional<at::Device> device,
::std::optional<bool> pin_memory) const;
// resize_ - in-place resize
const Tensor& resize_(
at::IntArrayRef size,
::std::optional<at::MemoryFormat> memory_format = ::std::nullopt) const;
// expand - expands tensor to new size
Tensor expand(at::IntArrayRef size, bool implicit = false) const;
// expand_as - expands to same size as another tensor
Tensor expand_as(const Tensor& other) const;
PaddleTensor _PD_GetInner() const { return tensor_; }
PaddleTensor& _PD_GetInner() { return tensor_; }
}; // NOLINT(readability/braces)
} // namespace at
@@ -0,0 +1,66 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#include <ATen/core/TensorBase.h>
#include <ATen/core/TensorBody.h>
#include <string_view>
namespace at {
void check_type(const TensorBase& tensor,
ScalarType type,
std::string_view type_name) {
PD_CHECK(tensor.scalar_type() == type,
"expected scalar type ",
type_name,
" but found ",
compat::_PD_AtenScalarTypeToPhiDataType(tensor.scalar_type()));
}
#define DEFINE_CAST(T, name) \
template <> \
PADDLE_API const T* TensorBase::const_data_ptr() const { \
check_type(*this, ScalarType::name, #name); \
return const_cast<T*>(tensor_.data<T>()); \
} \
\
template <> \
PADDLE_API const T* TensorBase::const_data_ptr<const T>() const { \
check_type(*this, ScalarType::name, #name); \
return const_cast<T*>(tensor_.data<std::remove_const_t<T>>()); \
} \
\
template <> \
PADDLE_API T* TensorBase::mutable_data_ptr() const { \
check_type(*this, ScalarType::name, #name); \
return const_cast<PaddleTensor&>(tensor_).data<T>(); \
} \
\
template <> \
PADDLE_API T* TensorBase::data_ptr() const { \
return const_cast<T*>(tensor_.data<T>()); \
}
AT_FORALL_SCALAR_TYPES_WITH_COMPLEX(DEFINE_CAST) // missing half and float16
// AT_FORALL_QINT_TYPES(DEFINE_CAST) // missing qint
DEFINE_CAST(uint16_t, UInt16)
DEFINE_CAST(uint32_t, UInt32)
DEFINE_CAST(uint64_t, UInt64)
#undef DEFINE_CAST
} // namespace at
@@ -0,0 +1,156 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#pragma once
#include <ostream>
#include <set>
#include <string>
#include <unordered_set>
#include <utility>
#include <vector>
namespace c10 {
/**
* class AliasInfo
*
* Data structure to hold aliasing information for an `Argument`. They can be
* nested to represent aliasing information on contained types.
*
* There is a `beforeSet` which describes the aliasing information before the
* operator executes, and an `afterSet` that describes aliasing info
* after execution.
*/
class AliasInfo {
public:
AliasInfo() = default;
AliasInfo(bool is_write,
const std::set<std::string>& before_qual_strings,
const std::set<std::string>& after_qual_strings)
: isWrite_(is_write) {
for (const auto& s : before_qual_strings) {
beforeSets_.insert(s);
}
for (const auto& s : after_qual_strings) {
afterSets_.insert(s);
}
}
bool isWrite() const { return isWrite_; }
const std::unordered_set<std::string>& beforeSets() const {
return beforeSets_;
}
const std::unordered_set<std::string>& afterSets() const {
return afterSets_;
}
// the alias info for the contained types of the type
// e.g. if this is an annotation on List[T], `sets` refers to
// the alias sets that the list may be in
// while containedTypes()[0] refers to the sets that members of the list
// may be in
void addContainedType(AliasInfo aliasInfo) {
containedTypes_.push_back(std::move(aliasInfo));
}
const std::vector<AliasInfo>& containedTypes() const {
return containedTypes_;
}
private:
std::unordered_set<std::string> beforeSets_;
std::unordered_set<std::string> afterSets_;
std::vector<AliasInfo> containedTypes_;
bool isWrite_ = false;
};
inline bool operator==(const AliasInfo& lhs, const AliasInfo& rhs) {
return lhs.isWrite() == rhs.isWrite() &&
lhs.beforeSets() == rhs.beforeSets() &&
lhs.afterSets() == rhs.afterSets() &&
lhs.containedTypes() == rhs.containedTypes();
}
// this does match the way things are represented in the schema
inline std::ostream& operator<<(std::ostream& out, const AliasInfo& aliasInfo) {
out << '(';
bool first = true;
for (const auto& set : aliasInfo.beforeSets()) {
if (first) {
first = false;
} else {
out << '|';
}
out << set;
}
if (aliasInfo.isWrite()) {
out << '!';
}
if (aliasInfo.beforeSets() != aliasInfo.afterSets()) {
out << " -> ";
first = true;
for (const auto& set : aliasInfo.afterSets()) {
if (first) {
first = false;
} else {
out << '|';
}
out << set;
}
}
out << ')';
return out;
}
} // namespace c10
inline std::size_t hash_combine(std::size_t lhs, std::size_t rhs) {
lhs ^= rhs + 0x9e3779b9 + (lhs << 6) + (lhs >> 2);
return lhs;
}
namespace std {
template <>
struct hash<c10::AliasInfo> {
size_t operator()(const c10::AliasInfo& aliasInfo) const {
auto hash = std::hash<bool>()(aliasInfo.isWrite());
// NOTE: for unordered_set hashes, we couldn't use hash_combine
// because hash_combine is order dependent. Instead, we choose to
// use XOR as the combining function as XOR is commutative.
size_t before_set_hash_seed = 0;
for (auto& e : aliasInfo.beforeSets()) {
auto symbol_hash = std::hash<std::string>()(e);
before_set_hash_seed = before_set_hash_seed ^ symbol_hash;
}
size_t after_set_hash_seed = 0;
for (auto& e : aliasInfo.afterSets()) {
auto symbol_hash = std::hash<std::string>()(e);
after_set_hash_seed = after_set_hash_seed ^ symbol_hash;
}
hash = hash_combine(hash, before_set_hash_seed);
hash = hash_combine(hash, after_set_hash_seed);
for (auto& e : aliasInfo.containedTypes()) {
auto contained_type_hash = std::hash<c10::AliasInfo>()(e);
hash = hash_combine(hash, contained_type_hash);
}
return hash;
}
};
} // namespace std
@@ -0,0 +1,201 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#include "ATen/core/function_schema.h"
namespace c10 {
namespace {
constexpr char kWildcardAliasSet[] = "*";
const char* schemaArgTypeName(SchemaArgType type) {
if (type == SchemaArgType::input) {
return "input";
}
if (type == SchemaArgType::output) {
return "output";
}
return "unknown";
}
bool aliasSetsMayOverlap(const std::unordered_set<std::string>& lhs,
const std::unordered_set<std::string>& rhs) {
if (lhs.empty() || rhs.empty()) {
return false;
}
if (lhs.count(kWildcardAliasSet) > 0 || rhs.count(kWildcardAliasSet) > 0) {
return true;
}
for (const auto& set : lhs) {
if (rhs.count(set) > 0) {
return true;
}
}
return false;
}
const Argument& getSchemaArgumentOrThrow(const FunctionSchema& schema,
const SchemaArgument& argument) {
const auto& args = schema.getCorrectList(argument);
TORCH_CHECK(argument.index < args.size(),
"Schema ",
schemaArgTypeName(argument.type),
" index ",
argument.index,
" is out of bounds for size ",
args.size());
return args.at(argument.index);
}
bool aliasInfoMayContainAlias(const AliasInfo& lhs,
const AliasInfo& rhs,
bool bidirectional) {
if (aliasSetsMayOverlap(lhs.afterSets(), rhs.afterSets())) {
return true;
}
for (const auto& child : lhs.containedTypes()) {
if (aliasInfoMayContainAlias(child, rhs, /*bidirectional=*/true)) {
return true;
}
}
if (!bidirectional) {
return false;
}
for (const auto& child : rhs.containedTypes()) {
if (aliasInfoMayContainAlias(lhs, child, /*bidirectional=*/true)) {
return true;
}
}
return false;
}
} // namespace
std::ostream& operator<<(std::ostream& out, const Argument& arg) {
out << arg.type()->str() << " " << arg.name();
if (arg.default_value()) {
out << " = " << arg.default_value();
}
return out;
}
std::ostream& operator<<(std::ostream& out, const FunctionSchema& schema) {
out << "(";
bool first = true;
for (const auto& arg : schema.arguments()) {
if (!first) {
out << ", ";
}
out << arg;
first = false;
}
if (schema.is_vararg()) {
if (!first) {
out << ", ";
}
out << "...";
}
out << ")";
out << " -> ";
if (schema.returns().size() == 1) {
out << schema.returns()[0];
} else {
out << "(";
first = true;
for (const auto& ret : schema.returns()) {
if (!first) {
out << ", ";
}
out << ret;
first = false;
}
out << ")";
}
return out;
}
std::optional<int> FunctionSchema::argumentIndexWithName(
const std::string& name) const {
for (size_t i = 0; i < arguments_.size(); ++i) {
if (arguments_[i].name() == name) {
return static_cast<int>(i);
}
}
return std::nullopt;
}
const std::vector<Argument>& FunctionSchema::getCorrectList(
const SchemaArgument& argument) const {
if (argument.type == SchemaArgType::input) {
return arguments();
}
if (argument.type == SchemaArgType::output) {
return returns();
}
TORCH_INTERNAL_ASSERT(false, "Could not match argument type");
}
bool FunctionSchema::is_aliasing(const SchemaArgument& argument) const {
const auto& arg = getSchemaArgumentOrThrow(*this, argument);
return arg.alias_info() != nullptr;
}
bool FunctionSchema::is_mutable(const SchemaArgument& argument) const {
const auto& arg = getSchemaArgumentOrThrow(*this, argument);
return arg.alias_info() != nullptr && arg.alias_info()->isWrite();
}
bool FunctionSchema::is_mutable(const std::string& name) const {
const auto index = argumentIndexWithName(name);
TORCH_CHECK(
index.has_value(), "Tried to test mutability of nonexistent name ", name);
return is_mutable({SchemaArgType::input, static_cast<size_t>(*index)});
}
bool FunctionSchema::may_alias(const SchemaArgument& lhs,
const SchemaArgument& rhs) const {
const auto& lhs_arg = getSchemaArgumentOrThrow(*this, lhs);
const auto& rhs_arg = getSchemaArgumentOrThrow(*this, rhs);
const auto* lhs_alias = lhs_arg.alias_info();
const auto* rhs_alias = rhs_arg.alias_info();
if (lhs_alias == nullptr || rhs_alias == nullptr) {
return false;
}
return aliasSetsMayOverlap(lhs_alias->afterSets(), rhs_alias->afterSets());
}
bool FunctionSchema::may_contain_alias(const SchemaArgument& lhs,
const SchemaArgument& rhs,
bool bidirectional) const {
const auto& lhs_arg = getSchemaArgumentOrThrow(*this, lhs);
const auto& rhs_arg = getSchemaArgumentOrThrow(*this, rhs);
const auto* lhs_alias = lhs_arg.alias_info();
const auto* rhs_alias = rhs_arg.alias_info();
if (lhs_alias == nullptr || rhs_alias == nullptr) {
return false;
}
return aliasInfoMayContainAlias(*lhs_alias, *rhs_alias, bidirectional);
}
} // namespace c10
@@ -0,0 +1,261 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#pragma once
#include <ATen/core/alias_info.h>
#include <ATen/core/ivalue.h>
#include <ATen/core/jit_type.h>
#include <string>
#include <vector>
#include "paddle/common/macros.h" // For macro PADDLE_API
namespace c10 {
struct Argument;
struct FunctionSchema;
enum class SchemaArgType;
struct SchemaArgument;
enum class SchemaArgType {
input,
output,
};
struct SchemaArgument {
SchemaArgType type;
size_t index;
};
struct PADDLE_API Argument {
Argument(std::string name = "",
const TypePtr& type = nullptr,
std::optional<int32_t> N = std::nullopt,
std::optional<torch::IValue> default_value = std::nullopt,
bool kwarg_only = false,
std::optional<c10::AliasInfo> alias_info = std::nullopt)
: Argument(std::move(name),
type,
type,
N,
std::move(default_value),
kwarg_only,
std::move(alias_info)) {}
Argument(std::string name,
TypePtr fake_type,
TypePtr real_type,
std::optional<int32_t> N = std::nullopt,
std::optional<torch::IValue> default_value = std::nullopt,
bool kwarg_only = false,
std::optional<c10::AliasInfo> alias_info = std::nullopt)
: name_(std::move(name)),
type_(fake_type ? std::move(fake_type) : TensorType::get()),
real_type_(real_type ? std::move(real_type) : type_),
N_(N),
default_value_(std::move(default_value)),
alias_info_(alias_info ? std::make_unique<c10::AliasInfo>(
std::move(*alias_info))
: nullptr),
kwarg_only_(kwarg_only) {
// this is an softly-enforced invariant for out arguments.
bool is_alias = alias_info_ != nullptr && alias_info_->isWrite();
is_out_ = kwarg_only_ && is_alias;
}
Argument(Argument&& rhs) noexcept = default;
Argument(const Argument& rhs)
: name_(rhs.name_),
type_(rhs.type_),
real_type_(rhs.real_type_),
N_(rhs.N_),
default_value_(rhs.default_value_),
alias_info_(rhs.alias_info_
? std::make_unique<c10::AliasInfo>(*rhs.alias_info_)
: nullptr),
kwarg_only_(rhs.kwarg_only_),
is_out_(rhs.is_out_) {}
Argument& operator=(Argument&& rhs) = default;
Argument& operator=(const Argument& rhs) {
if (this != &rhs) {
name_ = rhs.name_;
type_ = rhs.type_;
real_type_ = rhs.real_type_;
N_ = rhs.N_;
default_value_ = rhs.default_value_;
alias_info_ = rhs.alias_info_
? std::make_unique<c10::AliasInfo>(*rhs.alias_info_)
: nullptr;
kwarg_only_ = rhs.kwarg_only_;
is_out_ = rhs.is_out_;
}
return *this;
}
~Argument() = default;
const std::string& name() const { return name_; }
const TypePtr& type() const { return type_; }
// if type() is non-null, this is guaranteed to be non-null (if no real
// type was provided, this takes on type()'s value)
const TypePtr& real_type() const { return real_type_; }
const std::optional<int32_t>& N() const { return N_; }
const std::optional<torch::IValue>& default_value() const {
return default_value_;
}
bool kwarg_only() const { return kwarg_only_; }
bool is_out() const { return is_out_; }
[[nodiscard]] const c10::AliasInfo* alias_info() const {
return alias_info_.get();
}
bool is_inferred_type() const {
bool is_inferred_type = false;
TORCH_INTERNAL_ASSERT(type_);
if (auto pt = type_->cast<TensorType>()) {
if (pt->isInferredType()) {
is_inferred_type = true;
}
}
return is_inferred_type;
}
std::string formatTypeMismatchMsg(const std::string& actual_type) const {
std::string inferred_type_hint;
if (is_inferred_type()) {
inferred_type_hint = "Inferred '";
inferred_type_hint += name();
inferred_type_hint += "' to be of type 'Tensor' ";
inferred_type_hint +=
"because it was not annotated with an explicit type.\n";
}
std::string message;
message += "Expected a value of type '";
message += type()->repr_str();
message += "' for argument '";
message += name();
message += "' but instead found type '";
message += actual_type;
message += "'.\n";
message += inferred_type_hint;
return message;
}
Argument cloneWithType(const TypePtr& new_type) const {
return Argument(name_,
new_type,
N_,
default_value_,
kwarg_only_,
alias_info_ ? std::optional<c10::AliasInfo>(*alias_info_)
: std::nullopt);
}
friend PADDLE_API std::ostream& operator<<(std::ostream& out,
const Argument& arg);
private:
std::string name_;
TypePtr type_;
TypePtr real_type_; // this is ScalarType, not int, e.g.
// for list types, an optional statically known length for the list
// e.g. for int[3]: type = ListType::ofInts(), N = 3
// If present, this will allow scalars to be broadcast to this length to
// become a list.
std::optional<int32_t> N_;
std::optional<torch::IValue> default_value_;
// c10::AliasInfo is huge, so let's only allocate memory for it if
// necessary (which it isn't during schema parsing on startup, to
// give a pertinent example).
std::unique_ptr<c10::AliasInfo> alias_info_;
// is this only specifiable as a keyword argument?
bool kwarg_only_;
// marks if the argument is out variant of the schema
bool is_out_;
};
struct PADDLE_API FunctionSchema {
FunctionSchema(std::vector<Argument> arguments,
std::vector<Argument> returns,
bool is_vararg = false,
bool is_varret = false)
: arguments_(std::move(arguments)),
returns_(std::move(returns)),
is_vararg_(is_vararg),
is_varret_(is_varret) {
checkSchema();
}
const std::vector<Argument>& arguments() const { return arguments_; }
void checkSchema() const {
bool seen_default_arg = false;
for (const auto& arg : arguments()) {
if (arg.default_value()) {
seen_default_arg = true;
} else {
TORCH_INTERNAL_ASSERT(!seen_default_arg || arg.kwarg_only(),
"Non-default positional argument follows default "
"argument. Parameter ",
arg.name(),
" in ",
*this);
}
}
}
const std::vector<Argument>& returns() const { return returns_; }
bool is_vararg() const { return is_vararg_; }
bool is_varret() const { return is_varret_; }
std::optional<int> argumentIndexWithName(const std::string& name) const;
const std::vector<Argument>& getCorrectList(
const SchemaArgument& argument) const;
bool is_aliasing(const SchemaArgument& argument) const;
bool is_mutable(const SchemaArgument& argument) const;
bool is_mutable(const std::string& name) const;
bool may_alias(const SchemaArgument& lhs, const SchemaArgument& rhs) const;
bool may_contain_alias(const SchemaArgument& lhs,
const SchemaArgument& rhs,
bool bidirectional = true) const;
friend PADDLE_API std::ostream& operator<<(std::ostream& out,
const FunctionSchema& schema);
private:
std::vector<Argument> arguments_;
std::vector<Argument> returns_;
// if true then this schema takes an arbitrary number of additional arguments
// after the argument specified in arguments
// currently this is used primarily to represent 'primitive' operators whose
// arguments are not checked by schema
bool is_vararg_;
bool is_varret_;
};
PADDLE_API std::ostream& operator<<(std::ostream& out, const Argument& arg);
PADDLE_API std::ostream& operator<<(std::ostream& out,
const FunctionSchema& schema);
} // namespace c10
@@ -0,0 +1,72 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#pragma once
#include <c10/util/ArrayRef.h>
#include <vector>
namespace c10 {
// The passed in function must take T by value (T), or by
// const reference (const T&); taking T by non-const reference
// will result in an error like:
//
// error: no type named 'type' in 'class std::invoke_result<foobar::__lambda,
// T>'
//
// No explicit template parameters are required.
// Overload for explicit function and ArrayRef
template <class F, class T>
inline auto fmap(const T& inputs, const F& fn)
-> std::vector<decltype(fn(*inputs.begin()))> {
std::vector<decltype(fn(*inputs.begin()))> r;
r.reserve(inputs.size());
for (const auto& input : inputs) r.push_back(fn(input));
return r;
}
// C++ forbids taking an address of a constructor, so here's a workaround...
// Overload for constructor (R) application
template <typename R, typename T>
inline std::vector<R> fmap(const T& inputs) {
std::vector<R> r;
r.reserve(inputs.size());
for (auto& input : inputs) r.push_back(R(input));
return r;
}
template <typename F, typename T>
inline std::vector<T> filter(at::ArrayRef<T> inputs, const F& fn) {
std::vector<T> r;
r.reserve(inputs.size());
for (auto& input : inputs) {
if (fn(input)) {
r.push_back(input);
}
}
return r;
}
template <typename F, typename T>
inline std::vector<T> filter(const std::vector<T>& inputs, const F& fn) {
return filter<F, T>(static_cast<at::ArrayRef<T>>(inputs), fn);
}
} // namespace c10
@@ -0,0 +1,686 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#pragma once
#include <ATen/core/TensorBody.h>
#include <cstddef>
#include <iostream>
#include <memory>
#include <stdexcept>
#include <string>
#include <tuple>
#include <type_traits>
#include <utility>
#include <variant>
#include <vector>
namespace torch {
class CustomClassHolder {
public:
virtual ~CustomClassHolder() = default;
};
template <typename T>
class intrusive_ptr {
public:
using element_type = T;
using pointer = T*;
intrusive_ptr() : ptr_(nullptr) {}
intrusive_ptr(T* ptr) : ptr_(std::shared_ptr<T>(ptr)) {} // NOLINT
intrusive_ptr(std::shared_ptr<T> ptr) : ptr_(ptr) {} // NOLINT
template <typename... Args>
static intrusive_ptr<T> make(Args&&... args) {
return intrusive_ptr<T>(std::make_shared<T>(std::forward<Args>(args)...));
}
T* get() const { return ptr_.get(); }
T& operator*() const { return *ptr_; }
T* operator->() const { return ptr_.get(); }
// For IValue
std::shared_ptr<T> get_shared() const { return ptr_; }
explicit operator bool() const { return ptr_ != nullptr; }
private:
std::shared_ptr<T> ptr_;
};
template <typename T, typename... Args>
intrusive_ptr<T> make_intrusive(Args&&... args) {
return intrusive_ptr<T>::make(std::forward<Args>(args)...);
}
template <typename T>
struct _fake_type {};
enum class TypeTag {
None = 0,
Bool,
Int,
Double,
String,
Tensor,
GenericList,
CustomClass,
Tuple
};
class IValue; // Forward declaration
// Forward declaration of generic_to template function
template <typename T>
T generic_to(const IValue& ivalue, _fake_type<T>);
using GenericList = std::vector<IValue>;
// Separate tuple wrapper to avoid ambiguity with GenericList
struct GenericTuple {
std::vector<IValue> elements;
GenericTuple();
GenericTuple(std::vector<IValue> elems); // NOLINT
size_t size() const;
IValue& operator[](size_t idx);
const IValue& operator[](size_t idx) const;
};
class IValue {
private:
struct CustomClassWrapper {
std::shared_ptr<CustomClassHolder> ptr;
std::string class_name;
CustomClassWrapper(std::shared_ptr<CustomClassHolder> p,
const std::string& name)
: ptr(std::move(p)), class_name(name) {}
};
public:
IValue() : tag_(TypeTag::None), value_(std::monostate{}) {}
IValue(bool val) : tag_(TypeTag::Bool), value_(val) {} // NOLINT
IValue(int val) // NOLINT
: tag_(TypeTag::Int), value_(static_cast<int64_t>(val)) {}
IValue(int64_t val) : tag_(TypeTag::Int), value_(val) {} // NOLINT
IValue(double val) : tag_(TypeTag::Double), value_(val) {} // NOLINT
IValue(const std::string& val) // NOLINT
: tag_(TypeTag::String), value_(val) {}
IValue(std::string&& val) // NOLINT
: tag_(TypeTag::String), value_(std::move(val)) {}
IValue(const char* val) // NOLINT
: tag_(TypeTag::String), value_(std::string(val)) {}
IValue(at::Tensor val) : tag_(TypeTag::Tensor), value_(val) {} // NOLINT
IValue(ScalarType val) // NOLINT
: tag_(TypeTag::Int),
value_(static_cast<int64_t>(
static_cast<std::underlying_type_t<ScalarType>>(val))) {}
template <typename T>
IValue(intrusive_ptr<T> ptr) // NOLINT
: tag_(TypeTag::CustomClass),
value_(CustomClassWrapper{ptr.get_shared(), typeid(T).name()}) {}
template <typename T,
typename = std::enable_if_t<std::is_constructible_v<IValue, T>>>
IValue(const std::vector<T>& vec) // NOLINT
: tag_(TypeTag::GenericList) {
GenericList generic_list;
generic_list.reserve(vec.size());
for (const auto& item : vec) {
generic_list.emplace_back(IValue(item));
}
value_ = std::move(generic_list);
}
template <typename T,
typename = std::enable_if_t<std::is_constructible_v<IValue, T>>>
IValue(std::vector<T>&& vec) // NOLINT
: tag_(TypeTag::GenericList) {
GenericList generic_list;
generic_list.reserve(vec.size());
for (auto&& item : vec) {
generic_list.emplace_back(IValue(std::move(item)));
}
value_ = std::move(generic_list);
}
template <typename T,
typename = std::enable_if_t<std::is_constructible_v<IValue, T>>>
IValue(ArrayRef<T> arr) : IValue(arr.vec()) {} // NOLINT
template <typename T>
IValue(const std::optional<T>& opt) { // NOLINT
if (opt.has_value()) {
*this = IValue(*opt);
} else {
tag_ = TypeTag::None;
value_ = std::monostate{};
}
}
template <typename T>
IValue(std::optional<T>&& opt) { // NOLINT
if (opt.has_value()) {
*this = IValue(std::move(*opt));
} else {
tag_ = TypeTag::None;
value_ = std::monostate{};
}
}
// Variadic template constructor for tuple of any number of tensors or
// IValue-convertible types
template <typename... Args>
IValue(const std::tuple<Args...>& tuple_val) // NOLINT
: tag_(TypeTag::Tuple) {
static_assert(sizeof...(Args) > 0, "Tuple must have at least one element");
std::vector<IValue> elements;
elements.reserve(sizeof...(Args));
tuple_to_ivalue_vector(
tuple_val, elements, std::index_sequence_for<Args...>{});
value_ = GenericTuple(std::move(elements));
}
// Helper function to convert tuple elements to IValue vector using index
// sequence
template <typename Tuple, std::size_t... I>
void tuple_to_ivalue_vector(const Tuple& tuple_val,
std::vector<IValue>& elements, // NOLINT
std::index_sequence<I...>) {
(elements.emplace_back(std::get<I>(tuple_val)), ...);
}
IValue(const IValue& other) = default;
IValue(IValue&& other) = default;
IValue& operator=(const IValue& other) = default;
IValue& operator=(IValue&& other) = default;
bool is_none() const { return tag_ == TypeTag::None; }
bool is_bool() const { return tag_ == TypeTag::Bool; }
bool is_int() const { return tag_ == TypeTag::Int; }
bool is_double() const { return tag_ == TypeTag::Double; }
bool is_string() const { return tag_ == TypeTag::String; }
bool is_list() const { return tag_ == TypeTag::GenericList; }
bool is_tensor() const { return tag_ == TypeTag::Tensor; }
bool is_custom_class() const { return tag_ == TypeTag::CustomClass; }
bool is_tuple() const { return tag_ == TypeTag::Tuple; }
bool isNone() const { return is_none(); }
bool isBool() const { return is_bool(); }
bool isInt() const { return is_int(); }
bool isDouble() const { return is_double(); }
bool isString() const { return is_string(); }
bool isList() const { return is_list(); }
bool isTensor() const { return is_tensor(); }
bool isCustomClass() const { return is_custom_class(); }
bool isTuple() const { return is_tuple(); }
bool to_bool() const {
if (!is_bool()) throw std::runtime_error("Not a bool");
return std::get<bool>(value_);
}
int64_t to_int() const {
if (!is_int()) throw std::runtime_error("Not an int");
return std::get<int64_t>(value_);
}
double to_double() const {
if (!is_double()) throw std::runtime_error("Not a double");
return std::get<double>(value_);
}
const std::string& to_string() const {
if (!is_string()) throw std::runtime_error("Not a string");
return std::get<std::string>(value_);
}
const std::string_view to_string_view() const {
if (!is_string()) throw std::runtime_error("Not a string");
const auto& str = std::get<std::string>(value_);
return std::string_view(str.data(), str.size());
}
const GenericList& to_list() const {
if (!is_list()) throw std::runtime_error("Not a list");
return std::get<GenericList>(value_);
}
GenericList& to_list() {
if (!is_list()) throw std::runtime_error("Not a list");
return std::get<GenericList>(value_);
}
at::Tensor to_tensor() const {
if (!is_tensor()) throw std::runtime_error("Not a tensor");
return std::get<at::Tensor>(value_);
}
const GenericTuple& to_tuple() const {
if (!is_tuple()) throw std::runtime_error("Not a tuple");
return std::get<GenericTuple>(value_);
}
GenericTuple& to_tuple() {
if (!is_tuple()) throw std::runtime_error("Not a tuple");
return std::get<GenericTuple>(value_);
}
at::ScalarType to_scalar_type() const {
if (!is_int()) throw std::runtime_error("Not an int");
return static_cast<at::ScalarType>(std::get<int64_t>(value_));
}
bool toBool() const { return to_bool(); }
int64_t toInt() const { return to_int(); }
double toDouble() const { return to_double(); }
const std::string& toStringRef() const { return to_string(); }
std::string_view toStringView() const { return to_string_view(); }
at::Tensor toTensor() const { return to_tensor(); }
at::ScalarType toScalarType() const { return to_scalar_type(); }
std::string tagKind() const {
switch (tag_) {
case TypeTag::None:
return "None";
case TypeTag::Bool:
return "Bool";
case TypeTag::Int:
return "Int";
case TypeTag::Double:
return "Double";
case TypeTag::String:
return "String";
case TypeTag::Tensor:
return "Tensor";
case TypeTag::GenericList:
return "GenericList";
case TypeTag::CustomClass:
return "CustomClass";
case TypeTag::Tuple:
return "Tuple";
default:
return "InvalidTag";
}
}
template <typename T>
intrusive_ptr<T> to_custom_class() const {
if (!is_custom_class()) throw std::runtime_error("Not a custom class");
const auto& wrapper = std::get<CustomClassWrapper>(value_);
auto casted = std::dynamic_pointer_cast<T>(wrapper.ptr);
if (!casted) {
throw std::runtime_error("Cannot cast custom class to requested type");
}
return intrusive_ptr<T>(casted);
}
private:
template <typename T>
struct is_intrusive_ptr : std::false_type {};
template <typename T>
struct is_intrusive_ptr<intrusive_ptr<T>> : std::true_type {};
template <typename T>
static constexpr bool is_intrusive_ptr_v = is_intrusive_ptr<T>::value;
public:
bool try_to_bool(bool& out) const { // NOLINT
if (is_bool()) {
out = std::get<bool>(value_);
return true;
} else if (is_int()) {
out = (std::get<int64_t>(value_) != 0);
return true;
} else if (is_double()) {
out = (std::get<double>(value_) != 0.0);
return true;
}
return false;
}
bool try_to_int(int& out) const { // NOLINT
if (is_int()) {
out = static_cast<int>(std::get<int64_t>(value_));
return true;
} else if (is_double()) {
double val = std::get<double>(value_);
if (val != static_cast<int>(val)) {
std::cout << "Warning: Converting double(" << val
<< ") to int (precision loss)" << std::endl;
}
out = static_cast<int>(val);
return true;
}
return false;
}
bool try_to_double(double& out) const { // NOLINT
if (is_double()) {
out = std::get<double>(value_);
return true;
} else if (is_int()) {
out = static_cast<double>(std::get<int64_t>(value_));
return true;
}
return false;
}
bool try_to_string(std::string& out) const { // NOLINT
if (is_string()) {
out = std::get<std::string>(value_);
return true;
}
return false;
}
bool try_to_tensor(at::Tensor& out) const { // NOLINT
if (is_tensor()) {
out = std::get<at::Tensor>(value_);
return true;
}
return false;
}
bool try_to_scalar_type(at::ScalarType& out) const { // NOLINT
if (is_int()) {
out = static_cast<at::ScalarType>(std::get<int64_t>(value_));
return true;
}
return false;
}
template <typename T>
bool try_to_optional_type(std::optional<T>& out) const { // NOLINT
if (is_none()) {
out = std::nullopt;
return true;
} else {
T value;
if (try_convert_to<T>(value)) {
out = value;
return true;
}
}
return false;
}
bool try_to_custom_class(std::shared_ptr<CustomClassHolder>& out, // NOLINT
const std::string& expected_class_name) const {
if (is_custom_class()) {
const auto& wrapper = std::get<CustomClassWrapper>(value_);
if (wrapper.class_name == expected_class_name) {
out = wrapper.ptr;
return true;
}
}
return false;
}
template <typename T>
bool try_convert_to(T& out) const { // NOLINT
// Remove reference and cv-qualifiers from T
using BaseType = std::remove_cv_t<std::remove_reference_t<T>>;
if constexpr (std::is_same_v<BaseType, bool>) {
return try_to_bool(const_cast<bool&>(reinterpret_cast<const bool&>(out)));
} else if constexpr (std::is_same_v<BaseType, int>) {
return try_to_int(const_cast<int&>(reinterpret_cast<const int&>(out)));
} else if constexpr (std::is_same_v<BaseType, double>) {
return try_to_double(
const_cast<double&>(reinterpret_cast<const double&>(out)));
} else if constexpr (std::is_same_v<BaseType, std::string>) {
return try_to_string(
const_cast<std::string&>(reinterpret_cast<const std::string&>(out)));
} else if constexpr (std::is_same_v<BaseType, at::Tensor>) {
return try_to_tensor(
const_cast<at::Tensor&>(reinterpret_cast<const at::Tensor&>(out)));
} else if constexpr (std::is_same_v<BaseType, at::ScalarType>) {
return try_to_scalar_type(const_cast<at::ScalarType&>(
reinterpret_cast<const at::ScalarType&>(out)));
} else {
try {
// Handle const types by removing const and using const_cast
using NonConstType = std::remove_const_t<T>;
NonConstType temp = this->to<BaseType>();
const_cast<NonConstType&>(out) = std::move(temp);
return true;
} catch (const std::exception&) {
return false;
}
}
}
std::string get_custom_class_name() const {
if (!is_custom_class()) throw std::runtime_error("Not a custom class");
const auto& wrapper = std::get<CustomClassWrapper>(value_);
return wrapper.class_name;
}
template <typename T>
T to() && {
return generic_to(std::move(*this), _fake_type<T>{});
}
template <typename T>
T to() const& {
return generic_to(*this, _fake_type<T>{});
}
std::string type_string() const {
switch (tag_) {
case TypeTag::None:
return "None";
case TypeTag::Bool:
return "Bool";
case TypeTag::Int:
return "Int";
case TypeTag::Double:
return "Double";
case TypeTag::String:
return "String";
case TypeTag::Tensor:
return "Tensor";
case TypeTag::GenericList:
return "List";
case TypeTag::Tuple:
return "Tuple";
case TypeTag::CustomClass:
return "CustomClass(" + get_custom_class_name() + ")";
default:
return "Unknown";
}
}
std::string to_repr() const {
switch (tag_) {
case TypeTag::None:
return "None";
case TypeTag::Bool:
return std::get<bool>(value_) ? "true" : "false";
case TypeTag::Int:
return std::to_string(std::get<int64_t>(value_));
case TypeTag::Double:
return std::to_string(std::get<double>(value_));
case TypeTag::String:
return "\"" + std::get<std::string>(value_) + "\"";
case TypeTag::Tensor: {
const auto& tensor = std::get<at::Tensor>(value_);
return "Tensor(" + std::to_string(tensor.numel()) + " elements)";
}
case TypeTag::GenericList: {
const auto& list = std::get<GenericList>(value_);
std::string result = "[";
for (size_t i = 0; i < list.size(); ++i) {
if (i > 0) result += ", ";
result += list[i].to_repr();
}
result += "]";
return result;
}
case TypeTag::Tuple: {
const auto& tuple = std::get<GenericTuple>(value_);
std::string result = "(";
for (size_t i = 0; i < tuple.size(); ++i) {
if (i > 0) result += ", ";
result += tuple[i].to_repr();
}
if (tuple.size() == 1) result += ","; // Single element tuple
result += ")";
return result;
}
case TypeTag::CustomClass: {
const auto& wrapper = std::get<CustomClassWrapper>(value_);
return "CustomClass(" + wrapper.class_name + ")";
}
default:
return "Unknown";
}
}
friend std::ostream& operator<<(std::ostream& os, const IValue& val) {
return os << val.to_repr();
}
private:
TypeTag tag_;
std::variant<std::monostate,
bool,
int64_t,
double,
std::string,
at::Tensor,
GenericList,
CustomClassWrapper,
GenericTuple>
value_;
template <typename T>
friend T generic_to(const IValue& ivalue, _fake_type<T>);
};
inline GenericTuple::GenericTuple() = default;
inline GenericTuple::GenericTuple(std::vector<IValue> elems) // NOLINT
: elements(std::move(elems)) {}
inline size_t GenericTuple::size() const { return elements.size(); }
inline IValue& GenericTuple::operator[](size_t idx) { return elements[idx]; }
inline const IValue& GenericTuple::operator[](size_t idx) const {
return elements[idx];
}
template <>
inline bool generic_to(const IValue& ivalue, _fake_type<bool>) {
return ivalue.to_bool();
}
template <>
inline int generic_to(const IValue& ivalue, _fake_type<int>) {
return static_cast<int>(ivalue.to_int());
}
template <>
inline int64_t generic_to(const IValue& ivalue, _fake_type<int64_t>) {
return ivalue.to_int();
}
template <>
inline double generic_to(const IValue& ivalue, _fake_type<double>) {
return ivalue.to_double();
}
template <>
inline std::string generic_to(const IValue& ivalue, _fake_type<std::string>) {
return ivalue.to_string();
}
template <>
inline std::string_view generic_to(const IValue& ivalue,
_fake_type<std::string_view>) {
return ivalue.to_string_view();
}
template <>
inline at::Tensor generic_to(const IValue& ivalue, _fake_type<at::Tensor>) {
return ivalue.to_tensor();
}
template <typename T>
std::vector<T> generic_to(const IValue& ivalue, _fake_type<std::vector<T>>) {
auto list = ivalue.to_list();
std::vector<T> result;
result.reserve(list.size());
for (const auto& item : list) {
result.push_back(item.to<T>());
}
return result;
}
// Helper for converting IValue tuple to std::tuple using index sequence
template <typename Tuple, std::size_t... I>
Tuple ivalue_to_tuple_impl(const IValue& ivalue, std::index_sequence<I...>) {
const auto& generic_tuple = ivalue.to_tuple();
if (generic_tuple.size() != sizeof...(I)) {
throw std::runtime_error("Tuple size mismatch: expected " +
std::to_string(sizeof...(I)) + " but got " +
std::to_string(generic_tuple.size()));
}
// Use std::get<I> with index instead of type to avoid ambiguity
// when tuple contains multiple elements of the same type
return Tuple{generic_tuple[I].to<std::tuple_element_t<I, Tuple>>()...};
}
// Generic conversion from IValue to std::tuple
template <typename... Args>
std::tuple<Args...> generic_to(const IValue& ivalue,
_fake_type<std::tuple<Args...>>) {
return ivalue_to_tuple_impl<std::tuple<Args...>>(
ivalue, std::index_sequence_for<Args...>{});
}
template <typename T>
ArrayRef<T> generic_to(const IValue& ivalue, _fake_type<ArrayRef<T>>) {
static thread_local std::vector<T> temp_storage;
temp_storage = ivalue.to<std::vector<T>>();
return ArrayRef<T>(temp_storage);
}
template <typename T>
std::optional<T> generic_to(const IValue& ivalue,
_fake_type<std::optional<T>>) {
if (ivalue.is_none()) {
return std::nullopt;
}
return std::optional<T>(ivalue.to<T>());
}
template <typename T>
intrusive_ptr<T> generic_to(const IValue& ivalue,
_fake_type<intrusive_ptr<T>>) {
return ivalue.to_custom_class<T>();
}
} // namespace torch
namespace c10 {
using IValue = ::torch::IValue;
}
@@ -0,0 +1,385 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#pragma once
#include <ATen/core/jit_type_base.h>
#include <c10/util/Exception.h>
#include <memory>
#include <ostream>
#include <sstream>
#include <string>
#include <utility>
#include <vector>
namespace c10 {
inline bool operator!=(const Type& lhs, const Type& rhs) {
return !(lhs == rhs);
}
namespace detail {
// Lightweight runtime types used only by the compat schema parser.
class SchemaAtomicType final : public SharedType {
public:
static std::shared_ptr<SchemaAtomicType> create(TypeKind kind,
std::string repr) {
return std::shared_ptr<SchemaAtomicType>(
new SchemaAtomicType(kind, std::move(repr)));
}
bool equals(const Type& rhs) const override { return rhs.kind() == kind(); }
std::string str() const override { return repr_; }
private:
SchemaAtomicType(TypeKind kind, std::string repr)
: SharedType(kind), repr_(std::move(repr)) {}
std::string annotation_str_impl(
[[maybe_unused]] const TypePrinter& printer = nullptr) const override {
return repr_;
}
std::string repr_;
};
class SchemaOptionalType final : public SharedType {
public:
static const TypeKind Kind = TypeKind::OptionalType;
static std::shared_ptr<SchemaOptionalType> create(TypePtr elem) {
return std::shared_ptr<SchemaOptionalType>(
new SchemaOptionalType(std::move(elem)));
}
bool equals(const Type& rhs) const override {
if (rhs.kind() != kind()) {
return false;
}
const auto contained = rhs.containedTypes();
return contained.size() == 1 && *elem_.front() == *contained.front();
}
std::string str() const override { return elem_.front()->str() + "?"; }
at::ArrayRef<TypePtr> containedTypes() const override { return elem_; }
TypePtr createWithContained(
std::vector<TypePtr> contained_types) const override {
TORCH_CHECK(contained_types.size() == 1,
"Optional type expects exactly one contained type");
return create(std::move(contained_types.front()));
}
private:
explicit SchemaOptionalType(TypePtr elem)
: SharedType(Kind), elem_{std::move(elem)} {}
std::string annotation_str_impl(
const TypePrinter& printer = nullptr) const override {
return "Optional[" + elem_.front()->annotation_str(printer) + "]";
}
std::vector<TypePtr> elem_;
};
class SchemaTupleType final : public SharedType {
public:
static const TypeKind Kind = TypeKind::TupleType;
static std::shared_ptr<SchemaTupleType> create(
std::vector<TypePtr> elements) {
return std::shared_ptr<SchemaTupleType>(
new SchemaTupleType(std::move(elements)));
}
bool equals(const Type& rhs) const override {
if (rhs.kind() != kind()) {
return false;
}
const auto rhs_elems = rhs.containedTypes();
if (rhs_elems.size() != elements_.size()) {
return false;
}
for (size_t i = 0; i < elements_.size(); ++i) {
if (*elements_[i] != *rhs_elems[i]) {
return false;
}
}
return true;
}
std::string str() const override {
std::stringstream ss;
ss << "(";
for (size_t i = 0; i < elements_.size(); ++i) {
if (i > 0) {
ss << ", ";
}
ss << elements_[i]->str();
}
ss << ")";
return ss.str();
}
at::ArrayRef<TypePtr> containedTypes() const override { return elements_; }
TypePtr createWithContained(
std::vector<TypePtr> contained_types) const override {
return create(std::move(contained_types));
}
private:
explicit SchemaTupleType(std::vector<TypePtr> elements)
: SharedType(Kind), elements_(std::move(elements)) {}
std::string annotation_str_impl(
const TypePrinter& printer = nullptr) const override {
std::stringstream ss;
ss << "Tuple[";
for (size_t i = 0; i < elements_.size(); ++i) {
if (i > 0) {
ss << ", ";
}
ss << elements_[i]->annotation_str(printer);
}
ss << "]";
return ss.str();
}
std::vector<TypePtr> elements_;
};
} // namespace detail
inline TypePtr makeSchemaAtomicType(TypeKind kind, std::string repr) {
return detail::SchemaAtomicType::create(kind, std::move(repr));
}
inline TypePtr makeSchemaOptionalType(TypePtr elem) {
return detail::SchemaOptionalType::create(std::move(elem));
}
inline TypePtr makeSchemaTupleType(std::vector<TypePtr> elements) {
return detail::SchemaTupleType::create(std::move(elements));
}
struct TensorType;
using TensorTypePtr = SingletonTypePtr<TensorType>;
struct PADDLE_API TensorType : public Type {
bool equals(const Type& rhs) const override { return rhs.kind() == kind(); }
std::string str() const override { return "Tensor"; }
bool isInferredType() const { return is_inferred_; }
static const TypeKind Kind = TypeKind::TensorType;
static TensorTypePtr get() {
static TensorType value(/*inferred=*/false);
return TensorTypePtr(&value);
}
static TensorTypePtr getInferred() {
static TensorType value(/*inferred=*/true);
return TensorTypePtr(&value);
}
private:
explicit TensorType(bool inferred)
: Type(TypeKind::TensorType), is_inferred_(inferred) {}
bool is_inferred_;
};
struct NumberType;
using NumberTypePtr = SingletonTypePtr<NumberType>;
struct PADDLE_API NumberType : public Type {
bool equals(const Type& rhs) const override { return rhs.kind() == kind(); }
bool isSubtypeOfExt(const Type& rhs, std::ostream* why_not) const override {
return rhs.kind() == TypeKind::NumberType ||
Type::isSubtypeOfExt(rhs, why_not);
}
std::string str() const override { return "Scalar"; }
static const TypeKind Kind = TypeKind::NumberType;
static NumberTypePtr get() {
static NumberType value;
return NumberTypePtr(&value);
}
protected:
explicit NumberType(TypeKind kind = TypeKind::NumberType) : Type(kind) {}
std::string annotation_str_impl(
[[maybe_unused]] const TypePrinter& printer = nullptr) const override {
return "number";
}
};
struct FloatType;
using FloatTypePtr = SingletonTypePtr<FloatType>;
struct PADDLE_API FloatType : public NumberType {
bool equals(const Type& rhs) const override { return rhs.kind() == kind(); }
std::string str() const override { return "float"; }
bool isSubtypeOfExt(const Type& rhs, std::ostream* why_not) const override {
return rhs.kind() == TypeKind::NumberType ||
Type::isSubtypeOfExt(rhs, why_not);
}
static const TypeKind Kind = TypeKind::FloatType;
static FloatTypePtr get() {
static FloatType value;
return FloatTypePtr(&value);
}
private:
FloatType() : NumberType(TypeKind::FloatType) {}
std::string annotation_str_impl(
[[maybe_unused]] const TypePrinter& printer = nullptr) const override {
return "float";
}
};
struct IntType;
using IntTypePtr = SingletonTypePtr<IntType>;
struct PADDLE_API IntType : public NumberType {
bool equals(const Type& rhs) const override { return rhs.kind() == kind(); }
std::string str() const override { return "int"; }
bool isSubtypeOfExt(const Type& rhs, std::ostream* why_not) const override {
return rhs.kind() == TypeKind::NumberType ||
Type::isSubtypeOfExt(rhs, why_not);
}
static const TypeKind Kind = TypeKind::IntType;
static IntTypePtr get() {
static IntType value;
return IntTypePtr(&value);
}
private:
IntType() : NumberType(TypeKind::IntType) {}
std::string annotation_str_impl(
[[maybe_unused]] const TypePrinter& printer = nullptr) const override {
return "int";
}
};
struct BoolType;
using BoolTypePtr = SingletonTypePtr<BoolType>;
struct PADDLE_API BoolType : public Type {
bool equals(const Type& rhs) const override { return rhs.kind() == kind(); }
std::string str() const override { return "bool"; }
static const TypeKind Kind = TypeKind::BoolType;
static BoolTypePtr get() {
static BoolType value;
return BoolTypePtr(&value);
}
private:
BoolType() : Type(TypeKind::BoolType) {}
};
struct StringType;
using StringTypePtr = SingletonTypePtr<StringType>;
struct PADDLE_API StringType : public Type {
bool equals(const Type& rhs) const override { return rhs.kind() == kind(); }
std::string str() const override { return annotation_str(); }
std::string annotation_str_impl(
[[maybe_unused]] const TypePrinter& printer = nullptr) const override {
return "str";
}
static const TypeKind Kind = TypeKind::StringType;
static StringTypePtr get() {
static StringType value;
return StringTypePtr(&value);
}
private:
StringType() : Type(TypeKind::StringType) {}
};
struct NoneType;
using NoneTypePtr = SingletonTypePtr<NoneType>;
struct PADDLE_API NoneType : public Type {
bool equals(const Type& rhs) const override { return rhs.kind() == kind(); }
std::string str() const override { return "NoneType"; }
bool isSubtypeOfExt(const Type& rhs, std::ostream* why_not) const override {
return rhs.kind() == TypeKind::OptionalType ||
Type::isSubtypeOfExt(rhs, why_not);
}
static const TypeKind Kind = TypeKind::NoneType;
static NoneTypePtr get() {
static NoneType value;
return NoneTypePtr(&value);
}
private:
NoneType() : Type(TypeKind::NoneType) {}
};
struct DeviceObjType;
using DeviceObjTypePtr = SingletonTypePtr<DeviceObjType>;
struct PADDLE_API DeviceObjType : public Type {
bool equals(const Type& rhs) const override { return rhs.kind() == kind(); }
std::string str() const override { return "Device"; }
static const TypeKind Kind = TypeKind::DeviceObjType;
static DeviceObjTypePtr get() {
static DeviceObjType value;
return DeviceObjTypePtr(&value);
}
private:
DeviceObjType() : Type(TypeKind::DeviceObjType) {}
};
inline std::ostream& operator<<(std::ostream& out, const Type& t) {
out << t.str();
return out;
}
} // namespace c10
@@ -0,0 +1,337 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#pragma once
#include <functional>
#include <memory>
#include <optional>
#include <ostream>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>
#include "ATen/core/type_ptr.h"
#include "c10/util/ArrayRef.h"
#include "c10/util/Exception.h"
namespace c10 {
#define C10_FORALL_TYPES(_) \
_(TensorType) \
_(StringType) \
_(IntType) \
_(FloatType) \
_(BoolType) \
_(NoneType) \
_(TupleType) \
_(NumberType) \
_(OptionalType) \
_(UnionType) \
_(DeviceObjType) \
_(DynamicType)
enum class TypeKind {
#define DEFINE_TYPE(T) T,
C10_FORALL_TYPES(DEFINE_TYPE)
#undef DEFINE_TYPE
};
struct Type;
struct SharedType;
using TypePrinter = std::function<std::optional<std::string>(const Type&)>;
namespace detail {
template <typename T>
struct IsSingletonType : std::false_type {};
} // namespace detail
#define TORCH_DECLARE_SINGLETON(Type) \
struct Type; \
namespace detail { \
template <> \
struct IsSingletonType<Type> : std::true_type {}; \
}
TORCH_DECLARE_SINGLETON(NumberType)
TORCH_DECLARE_SINGLETON(TensorType)
TORCH_DECLARE_SINGLETON(StringType)
TORCH_DECLARE_SINGLETON(IntType)
TORCH_DECLARE_SINGLETON(FloatType)
TORCH_DECLARE_SINGLETON(BoolType)
TORCH_DECLARE_SINGLETON(NoneType)
TORCH_DECLARE_SINGLETON(TupleType)
TORCH_DECLARE_SINGLETON(OptionalType)
TORCH_DECLARE_SINGLETON(DeviceObjType)
namespace detail {
template <typename T, typename Enable = void>
struct CastReturnType {
using type = std::shared_ptr<T>;
};
template <typename T>
struct CastReturnType<T, std::enable_if_t<IsSingletonType<T>::value>> {
using type = SingletonTypePtr<T>;
};
template <typename T, typename Enable = void>
struct CastConstReturnType {
using type = std::shared_ptr<const T>;
};
template <typename T>
struct CastConstReturnType<T, std::enable_if_t<IsSingletonType<T>::value>> {
using type = SingletonTypePtr<const T>;
};
} // namespace detail
struct PADDLE_API Type {
friend PADDLE_API bool operator==(const Type& lhs, const Type& rhs);
private:
TypeKind kind_;
protected:
explicit Type(TypeKind kind) : kind_(kind) {}
Type(const Type&) = default;
Type& operator=(const Type&) = default;
Type(Type&&) noexcept = default;
Type& operator=(Type&&) noexcept = default;
virtual std::string annotation_str_impl(
const TypePrinter& /*printer*/) const {
return str();
}
virtual bool equals(const Type& rhs) const = 0;
virtual bool symmetric() const { return true; }
public:
template <typename T>
class SingletonOrSharedTypePtr {
public:
using element_type = typename std::shared_ptr<T>::element_type;
SingletonOrSharedTypePtr() = default;
SingletonOrSharedTypePtr(std::shared_ptr<T> x) // NOLINT(runtime/explicit)
: repr_(std::move(x)) {}
template <typename U,
std::enable_if_t<std::is_convertible_v<U*, T*>, bool> = true>
SingletonOrSharedTypePtr(std::shared_ptr<U> x) // NOLINT(runtime/explicit)
: repr_(std::move(x)) {}
SingletonOrSharedTypePtr(std::nullptr_t) // NOLINT(runtime/explicit)
: repr_(nullptr) {}
SingletonOrSharedTypePtr(SingletonTypePtr<T> p) // NOLINT(runtime/explicit)
: repr_(makeSingletonSharedPtr(p.get())) {}
template <typename U,
std::enable_if_t<std::is_convertible_v<U*, T*>, bool> = true>
SingletonOrSharedTypePtr(SingletonTypePtr<U> p) // NOLINT(runtime/explicit)
: repr_(makeSingletonSharedPtr(static_cast<T*>(p.get()))) {}
SingletonOrSharedTypePtr(const SingletonOrSharedTypePtr&) = default;
SingletonOrSharedTypePtr(SingletonOrSharedTypePtr&&) noexcept = default;
SingletonOrSharedTypePtr& operator=(const SingletonOrSharedTypePtr&) =
default;
SingletonOrSharedTypePtr& operator=(SingletonOrSharedTypePtr&&) noexcept =
default;
~SingletonOrSharedTypePtr() = default;
T* get() const { return repr_.get(); }
operator bool() const { return repr_ != nullptr; }
bool operator==(std::nullptr_t) const { return repr_ == nullptr; }
bool operator!=(std::nullptr_t) const { return repr_ != nullptr; }
template <typename U = T,
std::enable_if_t<!std::is_same_v<std::remove_const_t<U>, void>,
bool> = true>
U& operator*() const {
return *get();
}
T* operator->() const { return get(); }
private:
static std::shared_ptr<T> makeSingletonSharedPtr(T* ptr) {
return std::shared_ptr<T>(std::shared_ptr<T>(), ptr);
}
std::shared_ptr<T> repr_;
};
using TypePtr = SingletonOrSharedTypePtr<Type>;
virtual bool isSubtypeOfExt(const Type& rhs, std::ostream* why_not) const;
bool isSubtypeOf(const Type& rhs) const {
return isSubtypeOfExt(rhs, nullptr);
}
// Compatibility shims to accommodate existing code that passes shared_ptrs
// around. Ideally, we would just delete this, but it should be harmless.
template <typename T>
std::enable_if_t<std::is_base_of_v<Type, T>, bool> isSubtypeOf(
const std::shared_ptr<T>& rhs) const {
return isSubtypeOf(*rhs);
}
virtual std::string str() const = 0;
std::string annotation_str(const TypePrinter& printer) const {
if (printer) {
if (auto renamed = printer(*this)) {
return *renamed;
}
}
return annotation_str_impl(printer);
}
std::string annotation_str() const { return annotation_str(nullptr); }
virtual std::string repr_str() const { return annotation_str(); }
TypeKind kind() const { return kind_; }
template <typename T,
std::enable_if_t<!detail::IsSingletonType<T>::value, bool> = true>
typename detail::CastReturnType<T>::type cast() {
if (auto* typed = dynamic_cast<T*>(this)) {
return std::static_pointer_cast<T>(typed->shared_from_this());
}
return nullptr;
}
template <typename T,
std::enable_if_t<detail::IsSingletonType<T>::value, bool> = true>
typename detail::CastReturnType<T>::type cast() {
if (auto* typed = dynamic_cast<T*>(this)) {
return typename detail::CastReturnType<T>::type(typed);
}
return nullptr;
}
template <typename T,
std::enable_if_t<!detail::IsSingletonType<T>::value, bool> = true>
typename detail::CastConstReturnType<T>::type cast() const {
if (auto* typed = dynamic_cast<const T*>(this)) {
return std::static_pointer_cast<const T>(typed->shared_from_this());
}
return nullptr;
}
template <typename T,
std::enable_if_t<detail::IsSingletonType<T>::value, bool> = true>
typename detail::CastConstReturnType<T>::type cast() const {
if (auto* typed = dynamic_cast<const T*>(this)) {
return typename detail::CastConstReturnType<T>::type(typed);
}
return nullptr;
}
virtual ~Type() = default;
virtual at::ArrayRef<TypePtr> containedTypes() const { return {}; }
virtual TypePtr createWithContained(
std::vector<TypePtr> /*contained_types*/) const {
TORCH_CHECK(
false,
"type with contained types did not overload createWithContained: ",
str());
}
};
template <typename T>
using SingletonOrSharedTypePtr = Type::SingletonOrSharedTypePtr<T>;
template <typename T, typename U>
bool operator==(const SingletonOrSharedTypePtr<T>& x,
const SingletonOrSharedTypePtr<U>& y) {
return static_cast<const void*>(x.get()) == static_cast<const void*>(y.get());
}
template <typename T, typename U>
bool operator==(const SingletonOrSharedTypePtr<T>& x,
const SingletonTypePtr<U>& y) {
return static_cast<const void*>(x.get()) == static_cast<const void*>(y.get());
}
template <typename T, typename U>
bool operator==(const SingletonTypePtr<T>& x,
const SingletonOrSharedTypePtr<U>& y) {
return static_cast<const void*>(x.get()) == static_cast<const void*>(y.get());
}
template <typename T, typename U>
bool operator!=(const SingletonOrSharedTypePtr<T>& x,
const SingletonOrSharedTypePtr<U>& y) {
return !(x == y);
}
template <typename T, typename U>
bool operator!=(const SingletonOrSharedTypePtr<T>& x,
const SingletonTypePtr<U>& y) {
return !(x == y);
}
template <typename T, typename U>
bool operator!=(const SingletonTypePtr<T>& x,
const SingletonOrSharedTypePtr<U>& y) {
return !(x == y);
}
using TypePtr = SingletonOrSharedTypePtr<Type>;
// Base class for Types that are guaranteed to be owned by std::shared_ptr.
struct PADDLE_API SharedType : public Type,
public std::enable_shared_from_this<SharedType> {
using Type::Type;
};
inline bool operator==(const Type& lhs, const Type& rhs) {
if (!rhs.symmetric()) {
return rhs.equals(lhs);
}
return lhs.equals(rhs);
}
inline bool Type::isSubtypeOfExt(const Type& rhs, std::ostream* why_not) const {
if (*this == rhs) {
return true;
}
if (rhs.kind() == TypeKind::OptionalType) {
for (const auto& inner : rhs.containedTypes()) {
if (*this == *inner) {
return true;
}
}
}
return false;
}
} // namespace c10
namespace std {
template <typename T>
struct hash<c10::SingletonOrSharedTypePtr<T>> {
size_t operator()(const c10::SingletonOrSharedTypePtr<T>& x) const {
return std::hash<T*>()(x.get());
}
};
} // namespace std
@@ -0,0 +1,70 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#pragma once
#include <memory>
#include <type_traits>
#include "c10/util/Exception.h"
namespace c10 {
// Compatibility wrapper around a raw pointer so that existing code
// written to deal with a shared_ptr can keep working.
template <typename T>
class SingletonTypePtr {
public:
SingletonTypePtr(T* p) : repr_(p) {} // NOLINT(runtime/explicit)
// We need this to satisfy Pybind11, but it shouldn't be hit.
explicit SingletonTypePtr(std::shared_ptr<T> /*unused*/) {
TORCH_CHECK(false);
}
using element_type = typename std::shared_ptr<T>::element_type;
template <typename U = T,
std::enable_if_t<!std::is_same_v<std::remove_const_t<U>, void>,
bool> = true>
T& operator*() const {
return *repr_;
}
T* get() const { return repr_; }
T* operator->() const { return repr_; }
operator bool() const { return repr_ != nullptr; }
private:
T* repr_{nullptr};
};
template <typename T, typename U>
bool operator==(SingletonTypePtr<T> lhs, SingletonTypePtr<U> rhs) {
return static_cast<const void*>(lhs.get()) ==
static_cast<const void*>(rhs.get());
}
template <typename T, typename U>
bool operator!=(SingletonTypePtr<T> lhs, SingletonTypePtr<U> rhs) {
return !(lhs == rhs);
}
} // namespace c10
@@ -0,0 +1,201 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#if defined(PADDLE_WITH_CUDA)
#include <ATen/cuda/CUDABlas.h>
#include "paddle/phi/backends/dynload/cublas.h"
#include "paddle/phi/core/enforce.h"
namespace at::cuda::blas {
namespace {
inline cublasOperation_t to_cublas_op(char trans) {
switch (trans) {
case 'T':
case 't':
return CUBLAS_OP_T;
case 'N':
case 'n':
return CUBLAS_OP_N;
case 'C':
case 'c':
return CUBLAS_OP_C;
default:
PADDLE_THROW(common::errors::InvalidArgument(
"at::cuda::blas::gemm: invalid transpose character '%c'", trans));
}
}
} // namespace
/* ───────────── gemm<double> ───────────── */
template <>
void gemm<double>(CUDABLAS_GEMM_ARGTYPES(double)) {
cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
cublasOperation_t opa = to_cublas_op(transa);
cublasOperation_t opb = to_cublas_op(transb);
PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cublasDgemm(handle,
opa,
opb,
static_cast<int>(m),
static_cast<int>(n),
static_cast<int>(k),
&alpha,
a,
static_cast<int>(lda),
b,
static_cast<int>(ldb),
&beta,
c,
static_cast<int>(ldc)));
}
/* ───────────── gemm<float> ───────────── */
template <>
void gemm<float>(CUDABLAS_GEMM_ARGTYPES(float)) {
cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
cublasOperation_t opa = to_cublas_op(transa);
cublasOperation_t opb = to_cublas_op(transb);
PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cublasSgemm(handle,
opa,
opb,
static_cast<int>(m),
static_cast<int>(n),
static_cast<int>(k),
&alpha,
a,
static_cast<int>(lda),
b,
static_cast<int>(ldb),
&beta,
c,
static_cast<int>(ldc)));
}
/* ───────────── gemm<c10::complex<double>> ───────────── */
template <>
void gemm<c10::complex<double>>(CUDABLAS_GEMM_ARGTYPES(c10::complex<double>)) {
cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
cublasOperation_t opa = to_cublas_op(transa);
cublasOperation_t opb = to_cublas_op(transb);
PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cublasZgemm(
handle,
opa,
opb,
static_cast<int>(m),
static_cast<int>(n),
static_cast<int>(k),
reinterpret_cast<const cuDoubleComplex *>(&alpha),
reinterpret_cast<const cuDoubleComplex *>(a),
static_cast<int>(lda),
reinterpret_cast<const cuDoubleComplex *>(b),
static_cast<int>(ldb),
reinterpret_cast<const cuDoubleComplex *>(&beta),
reinterpret_cast<cuDoubleComplex *>(c),
static_cast<int>(ldc)));
}
/* ───────────── gemm<c10::complex<float>> ───────────── */
template <>
void gemm<c10::complex<float>>(CUDABLAS_GEMM_ARGTYPES(c10::complex<float>)) {
cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
cublasOperation_t opa = to_cublas_op(transa);
cublasOperation_t opb = to_cublas_op(transb);
PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cublasCgemm(
handle,
opa,
opb,
static_cast<int>(m),
static_cast<int>(n),
static_cast<int>(k),
reinterpret_cast<const cuFloatComplex *>(&alpha),
reinterpret_cast<const cuFloatComplex *>(a),
static_cast<int>(lda),
reinterpret_cast<const cuFloatComplex *>(b),
static_cast<int>(ldb),
reinterpret_cast<const cuFloatComplex *>(&beta),
reinterpret_cast<cuFloatComplex *>(c),
static_cast<int>(ldc)));
}
/* ───────────── gemm<at::Half> ───────────── */
template <>
void gemm<at::Half>(CUDABLAS_GEMM_ARGTYPES(at::Half)) {
cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
cublasOperation_t opa = to_cublas_op(transa);
cublasOperation_t opb = to_cublas_op(transb);
// Use cublasGemmEx with FP32 compute for Half inputs
float alpha_f = alpha;
float beta_f = beta;
PADDLE_ENFORCE_GPU_SUCCESS(
phi::dynload::cublasGemmEx(handle,
opa,
opb,
static_cast<int>(m),
static_cast<int>(n),
static_cast<int>(k),
&alpha_f,
a,
CUDA_R_16F,
static_cast<int>(lda),
b,
CUDA_R_16F,
static_cast<int>(ldb),
&beta_f,
c,
CUDA_R_16F,
static_cast<int>(ldc),
CUDA_R_32F,
CUBLAS_GEMM_DEFAULT_TENSOR_OP));
}
/* ───────────── gemm<at::BFloat16> ───────────── */
template <>
void gemm<at::BFloat16>(CUDABLAS_GEMM_ARGTYPES(at::BFloat16)) {
cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
cublasOperation_t opa = to_cublas_op(transa);
cublasOperation_t opb = to_cublas_op(transb);
// Use cublasGemmEx with FP32 compute for BFloat16 inputs
float alpha_f = alpha;
float beta_f = beta;
PADDLE_ENFORCE_GPU_SUCCESS(
phi::dynload::cublasGemmEx(handle,
opa,
opb,
static_cast<int>(m),
static_cast<int>(n),
static_cast<int>(k),
&alpha_f,
a,
CUDA_R_16BF,
static_cast<int>(lda),
b,
CUDA_R_16BF,
static_cast<int>(ldb),
&beta_f,
c,
CUDA_R_16BF,
static_cast<int>(ldc),
CUDA_R_32F,
CUBLAS_GEMM_DEFAULT_TENSOR_OP));
}
} // namespace at::cuda::blas
#endif // PADDLE_WITH_CUDA
@@ -0,0 +1,73 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#pragma once
/*
Provides a subset of CUDA BLAS functions as templates:
gemm<Dtype>(transa, transb, m, n, k, alpha, a, lda, b, ldb, beta, c,
ldc)
where Dtype is double, float, c10::complex<double>, c10::complex<float>,
at::Half or at::BFloat16. The functions are available in at::cuda::blas
namespace.
*/
#include <ATen/OpMathType.h>
#include <ATen/cuda/CUDAContext.h>
#include "paddle/common/macros.h"
namespace at::cuda::blas {
/* LEVEL 3 BLAS FUNCTIONS */
#define CUDABLAS_GEMM_ARGTYPES(Dtype) \
CUDABLAS_GEMM_ARGTYPES_AND_C_DTYPE(Dtype, Dtype)
#define CUDABLAS_GEMM_ARGTYPES_AND_C_DTYPE(Dtype, C_Dtype) \
char transa, char transb, int64_t m, int64_t n, int64_t k, \
at::opmath_type<Dtype> alpha, const Dtype *a, int64_t lda, \
const Dtype *b, int64_t ldb, at::opmath_type<Dtype> beta, C_Dtype *c, \
int64_t ldc
#define CUDABLAS_GEMM_ARGS(Dtype) \
transa, transb, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc
template <typename Dtype, typename C_Dtype = Dtype>
inline void gemm(CUDABLAS_GEMM_ARGTYPES_AND_C_DTYPE(Dtype, C_Dtype)) {
static_assert(false && sizeof(Dtype),
"at::cuda::blas::gemm: not implemented");
}
template <>
PADDLE_API void gemm<double>(CUDABLAS_GEMM_ARGTYPES(double));
template <>
PADDLE_API void gemm<float>(CUDABLAS_GEMM_ARGTYPES(float));
template <>
PADDLE_API void gemm<c10::complex<double>>(
CUDABLAS_GEMM_ARGTYPES(c10::complex<double>));
template <>
PADDLE_API void gemm<c10::complex<float>>(
CUDABLAS_GEMM_ARGTYPES(c10::complex<float>));
template <>
PADDLE_API void gemm<at::Half>(CUDABLAS_GEMM_ARGTYPES(at::Half));
template <>
PADDLE_API void gemm<at::BFloat16>(CUDABLAS_GEMM_ARGTYPES(at::BFloat16));
} // namespace at::cuda::blas
@@ -0,0 +1,214 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
#include <ATen/cuda/CUDAContext.h>
#include <c10/core/Allocator.h>
#include <mutex>
#include "paddle/phi/backends/context_pool.h"
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/backends/gpu/gpu_info.h"
#include "paddle/phi/core/memory/allocation/allocator_facade.h"
namespace at::cuda {
namespace {
inline void ensureDeviceContextPoolInitialized() {
static std::once_flag init_pool_once;
std::call_once(init_pool_once, []() {
if (phi::DeviceContextPool::IsInitialized()) {
return;
}
std::vector<phi::Place> places;
int gpu_count = phi::backends::gpu::GetGPUDeviceCount();
for (int device = 0; device < gpu_count; ++device) {
places.emplace_back(phi::GPUPlace(device));
}
places.emplace_back(phi::CPUPlace());
places.emplace_back(phi::GPUPinnedPlace());
phi::DeviceContextPool::Init(places);
});
}
/// Returns the GPUContext for the current device.
inline phi::GPUContext* getCurrentGPUContext() {
ensureDeviceContextPoolInitialized();
int device_id = phi::backends::gpu::GetCurrentDeviceId();
return static_cast<phi::GPUContext*>(
phi::DeviceContextPool::Instance().Get(phi::GPUPlace(device_id)));
}
/// Frees a phi::Allocation that was released with .release() during allocate().
static void deletePaddleCUDAAllocation(void* p) {
delete static_cast<phi::Allocation*>(p);
}
/// Adapter class that wraps Paddle's AllocatorFacade as a c10::Allocator.
/// This provides a bridge between Paddle's allocation interface and PyTorch's
/// c10::Allocator interface for the CUDA compatibility layer.
class PaddleCUDAAllocatorAdapter : public c10::Allocator {
public:
c10::DataPtr allocate(size_t n) override {
int device_id = phi::backends::gpu::GetCurrentDeviceId();
if (n == 0) {
// Return a DataPtr that carries the current CUDA device without
// allocating any memory. Callers that probe device identity via
// DataPtr::device() (e.g. zero-byte tensor construction) will therefore
// observe the correct CUDA device rather than a default CPU device.
// NOTE: For HIP/ROCm builds, PyTorch's compatibility layer still
// exposes DeviceType::CUDA (kCUDA) rather than a separate HIP device
// type, so we follow the same convention here.
return c10::DataPtr(nullptr,
nullptr,
nullptr,
c10::Device(c10::DeviceType::CUDA, device_id));
}
auto* alloc = paddle::memory::allocation::AllocatorFacade::Instance()
.GetAllocator(phi::GPUPlace(device_id))
.get();
auto phi_alloc = alloc->Allocate(n);
void* ptr = phi_alloc->ptr();
phi::Place place = phi_alloc->place();
// Transfer ownership of phi_alloc to the DataPtr's context.
auto* raw_alloc = phi_alloc.release();
return c10::DataPtr(
ptr, raw_alloc, deletePaddleCUDAAllocation, c10::Device(place));
}
void copy_data(void* dst, const void* src, size_t n) const override {
if (n == 0) return;
// Use GPU device-to-device copy. std::memcpy is not valid for device
// memory; callers such as c10::Allocator::clone() rely on this method to
// perform correct D2D copies on CUDA/HIP memory.
#ifdef PADDLE_WITH_HIP
PADDLE_ENFORCE_GPU_SUCCESS(hipMemcpy(dst, src, n, hipMemcpyDeviceToDevice));
#else
PADDLE_ENFORCE_GPU_SUCCESS(
cudaMemcpy(dst, src, n, cudaMemcpyDeviceToDevice));
#endif
}
c10::DeleterFnPtr raw_deleter() const override {
// allocate() returns data=device_ptr, context=phi::Allocation*, so
// get() != get_context() and the raw_allocate/raw_deallocate API is
// unsafe for this allocator. Returning nullptr signals that.
return nullptr;
}
};
} // namespace
CUDAContextDeviceProp* getCurrentDeviceProperties() {
int device = phi::backends::gpu::GetCurrentDeviceId();
return getDeviceProperties(device);
}
int warp_size() { return getCurrentDeviceProperties()->warpSize; }
CUDAContextDeviceProp* getDeviceProperties(c10::DeviceIndex device) {
return const_cast<CUDAContextDeviceProp*>(
&phi::backends::gpu::GetDeviceProperties(device));
}
bool canDeviceAccessPeer(c10::DeviceIndex device,
c10::DeviceIndex peer_device) {
int can_access = 0;
#ifdef PADDLE_WITH_HIP
hipDeviceCanAccessPeer(&can_access, device, peer_device);
#else
cudaDeviceCanAccessPeer(&can_access, device, peer_device);
#endif
return can_access != 0;
}
/* Handles */
CUDAContextSparseHandle getCurrentCUDASparseHandle() {
return getCurrentGPUContext()->cusparse_handle();
}
CUDAContextBlasHandle getCurrentCUDABlasHandle() {
return getCurrentGPUContext()->cublas_handle();
}
CUDAContextBlasLtHandle getCurrentCUDABlasLtHandle() {
return getCurrentGPUContext()->cublaslt_handle();
}
void clearCublasWorkspaces() {
// Workspaces are owned and managed by phi::GPUContext; no explicit
// cleanup is required here.
}
WorkspaceMapWithMutex& cublas_handle_stream_to_workspace() {
static WorkspaceMapWithMutex workspace_map;
return workspace_map;
}
WorkspaceMapWithMutex& cublaslt_handle_stream_to_workspace() {
static WorkspaceMapWithMutex workspace_map;
return workspace_map;
}
// Default workspace size consistent with PyTorch's chosen default (32 MiB).
static constexpr size_t kDefaultWorkspaceSize = 32UL * 1024UL * 1024UL;
size_t getChosenWorkspaceSize() { return kDefaultWorkspaceSize; }
size_t getCUDABlasLtWorkspaceSize() {
// Probe the context with the default size and return what was actually
// allocated.
auto [ptr, size] =
getCurrentGPUContext()->cublaslt_workspace(kDefaultWorkspaceSize);
(void)ptr;
return size;
}
void* getCUDABlasLtWorkspace() {
return getCurrentGPUContext()
->cublaslt_workspace(kDefaultWorkspaceSize)
.first;
}
CUDAContextSolverHandle getCurrentCUDASolverDnHandle() {
return getCurrentGPUContext()->cusolver_dn_handle();
}
#if defined(USE_CUDSS)
cudssHandle_t getCurrentCudssHandle() {
// cudss is not yet integrated into phi::GPUContext; not implemented.
PADDLE_THROW(
common::errors::Unimplemented("getCurrentCudssHandle() is not "
"implemented in the Paddle compat layer."));
return nullptr;
}
#endif // USE_CUDSS
c10::Allocator* getCUDADeviceAllocator() {
static PaddleCUDAAllocatorAdapter adapter;
return &adapter;
}
} // namespace at::cuda
#endif // PADDLE_WITH_CUDA || PADDLE_WITH_HIP
@@ -0,0 +1,26 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#pragma once
#include <ATen/cuda/CUDAContextLight.h>
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
#include <ATen/cuda/Exceptions.h>
#include <c10/cuda/CUDAStream.h>
#endif
@@ -0,0 +1,136 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#pragma once
// Light-weight version of CUDAContext.h with fewer transitive includes
// cublasLT was introduced in CUDA 10.1 but we enable only for 11.1 that also
// added bf16 support
#if defined(PADDLE_WITH_HIP)
#include <hip/hip_runtime.h>
#elif defined(PADDLE_WITH_CUDA)
#if defined(USE_CUDSS)
#include <cudss.h>
#endif
#include <driver_types.h>
#endif
#include <c10/core/Allocator.h>
#include <c10/cuda/CUDAFunctions.h>
#include <cstdint>
#include <map>
#include <shared_mutex>
#include <tuple>
#include "paddle/common/macros.h"
#include "paddle/phi/backends/gpu/forwards.h"
namespace c10 {
struct Allocator;
}
namespace at::cuda {
#if defined(PADDLE_WITH_HIP)
using CUDAContextDeviceProp = phi::gpuDeviceProp;
using CUDAContextSparseHandle = phi::sparseHandle_t;
using CUDAContextBlasHandle = phi::blasHandle_t;
using CUDAContextBlasLtHandle = phi::blasLtHandle_t;
using CUDAContextSolverHandle = phi::solverHandle_t;
#elif defined(PADDLE_WITH_CUDA)
using CUDAContextDeviceProp = cudaDeviceProp;
using CUDAContextSparseHandle = cusparseHandle_t;
using CUDAContextBlasHandle = cublasHandle_t;
using CUDAContextBlasLtHandle = cublasLtHandle_t;
using CUDAContextSolverHandle = cusolverDnHandle_t;
#endif
/*
A common CUDA interface for ATen.
This interface is distinct from CUDAHooks, which defines an interface that links
to both CPU-only and CUDA builds. That interface is intended for runtime
dispatch and should be used from files that are included in both CPU-only and
CUDA builds.
CUDAContext, on the other hand, should be preferred by files only included in
CUDA builds. It is intended to expose CUDA functionality in a consistent
manner.
This means there is some overlap between the CUDAContext and CUDAHooks, but
the choice of which to use is simple: use CUDAContext when in a CUDA-only file,
use CUDAHooks otherwise.
Note that CUDAContext simply defines an interface with no associated class.
It is expected that the modules whose functions compose this interface will
manage their own state. There is only a single CUDA context/state.
*/
/**
* DEPRECATED: use device_count() instead
*/
inline int64_t getNumGPUs() { return c10::cuda::device_count(); }
/**
* CUDA is available if we compiled with CUDA, and there are one or more
* devices. If we compiled with CUDA but there is a driver problem, etc.,
* this function will report CUDA is not available (rather than raise an error.)
*/
inline bool is_available() { return c10::cuda::device_count() > 0; }
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
PADDLE_API CUDAContextDeviceProp* getCurrentDeviceProperties();
PADDLE_API int warp_size();
PADDLE_API CUDAContextDeviceProp* getDeviceProperties(c10::DeviceIndex device);
PADDLE_API bool canDeviceAccessPeer(c10::DeviceIndex device,
c10::DeviceIndex peer_device);
/* Handles */
PADDLE_API CUDAContextSparseHandle getCurrentCUDASparseHandle();
PADDLE_API CUDAContextBlasHandle getCurrentCUDABlasHandle();
PADDLE_API CUDAContextBlasLtHandle getCurrentCUDABlasLtHandle();
PADDLE_API void clearCublasWorkspaces();
struct WorkspaceMapWithMutex {
std::map<std::tuple<void*, void*>, at::DataPtr> map;
std::shared_mutex mutex;
};
PADDLE_API WorkspaceMapWithMutex& cublas_handle_stream_to_workspace();
PADDLE_API WorkspaceMapWithMutex& cublaslt_handle_stream_to_workspace();
PADDLE_API size_t getChosenWorkspaceSize();
PADDLE_API size_t getCUDABlasLtWorkspaceSize();
PADDLE_API void* getCUDABlasLtWorkspace();
PADDLE_API CUDAContextSolverHandle getCurrentCUDASolverDnHandle();
#if defined(USE_CUDSS)
PADDLE_API cudssHandle_t getCurrentCudssHandle();
#endif
// Get the CUDA device allocator for the current device.
// Returns a pointer to a c10::Allocator that allocates GPU memory.
PADDLE_API c10::Allocator* getCUDADeviceAllocator();
#endif
} // namespace at::cuda
@@ -0,0 +1,164 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#pragma once
#include <c10/core/ScalarType.h>
#if defined(PADDLE_WITH_HIP)
#include <hip/hip_runtime.h>
#include <hip/library_types.h>
#elif defined(PADDLE_WITH_CUDA)
#include <cuda.h>
#include <library_types.h>
#endif
namespace at::cuda {
#if defined(PADDLE_WITH_HIP)
using cudaDataType = hipDataType;
#define CUDA_R_16F HIP_R_16F
#define CUDA_R_32F HIP_R_32F
#define CUDA_R_64F HIP_R_64F
#define CUDA_C_16F HIP_C_16F
#define CUDA_C_32F HIP_C_32F
#define CUDA_C_64F HIP_C_64F
#define CUDA_R_8U HIP_R_8U
#define CUDA_R_8I HIP_R_8I
#define CUDA_R_32I HIP_R_32I
#define CUDA_R_16I HIP_R_16I
#define CUDA_R_64I HIP_R_64I
#define CUDA_R_16BF HIP_R_16BF
#define CUDA_R_8F_E4M3 HIP_R_8F_E4M3
#define CUDA_R_8F_E5M2 HIP_R_8F_E5M2
#elif defined(PADDLE_WITH_CUDA)
using cudaDataType = cudaDataType;
#endif
template <typename scalar_t>
cudaDataType getCudaDataType() {
static_assert(false && sizeof(scalar_t),
"Cannot convert type to cudaDataType.");
return {};
}
template <>
inline cudaDataType getCudaDataType<at::Half>() {
return CUDA_R_16F;
}
template <>
inline cudaDataType getCudaDataType<float>() {
return CUDA_R_32F;
}
template <>
inline cudaDataType getCudaDataType<double>() {
return CUDA_R_64F;
}
template <>
inline cudaDataType getCudaDataType<c10::complex<c10::Half>>() {
return CUDA_C_16F;
}
template <>
inline cudaDataType getCudaDataType<c10::complex<float>>() {
return CUDA_C_32F;
}
template <>
inline cudaDataType getCudaDataType<c10::complex<double>>() {
return CUDA_C_64F;
}
template <>
inline cudaDataType getCudaDataType<uint8_t>() {
return CUDA_R_8U;
}
template <>
inline cudaDataType getCudaDataType<int8_t>() {
return CUDA_R_8I;
}
template <>
inline cudaDataType getCudaDataType<int>() {
return CUDA_R_32I;
}
template <>
inline cudaDataType getCudaDataType<int16_t>() {
return CUDA_R_16I;
}
template <>
inline cudaDataType getCudaDataType<int64_t>() {
return CUDA_R_64I;
}
template <>
inline cudaDataType getCudaDataType<at::BFloat16>() {
return CUDA_R_16BF;
}
inline cudaDataType ScalarTypeToCudaDataType(
const c10::ScalarType& scalar_type) {
switch (scalar_type) {
case c10::ScalarType::Byte:
return CUDA_R_8U;
case c10::ScalarType::Char:
return CUDA_R_8I;
case c10::ScalarType::Int:
return CUDA_R_32I;
case c10::ScalarType::Half:
return CUDA_R_16F;
case c10::ScalarType::Float:
return CUDA_R_32F;
case c10::ScalarType::Double:
return CUDA_R_64F;
// case c10::ScalarType::ComplexHalf:
// return CUDA_C_16F;
case c10::ScalarType::ComplexFloat:
return CUDA_C_32F;
case c10::ScalarType::ComplexDouble:
return CUDA_C_64F;
case c10::ScalarType::Short:
return CUDA_R_16I;
case c10::ScalarType::Long:
return CUDA_R_64I;
case c10::ScalarType::BFloat16:
return CUDA_R_16BF;
#if defined(PADDLE_WITH_HIP)
case c10::ScalarType::Float8_e4m3fn:
return CUDA_R_8F_E4M3;
case c10::ScalarType::Float8_e5m2:
return CUDA_R_8F_E5M2;
case c10::ScalarType::Float8_e4m3fnuz:
return HIP_R_8F_E4M3_FNUZ;
case c10::ScalarType::Float8_e5m2fnuz:
return HIP_R_8F_E5M2_FNUZ;
#elif !defined(USE_ROCM) || ROCM_VERSION >= 60300
case c10::ScalarType::Float8_e4m3fn:
return CUDA_R_8F_E4M3;
case c10::ScalarType::Float8_e5m2:
return CUDA_R_8F_E5M2;
#endif
// #if (defined(CUDA_VERSION) && CUDA_VERSION >= 12080) ||
// (defined(USE_ROCM) && ROCM_VERSION >= 70000)
// case c10::ScalarType::Float4_e2m1fn_x2:
// return CUDA_R_4F_E2M1;
// #endif
default:
TORCH_INTERNAL_ASSERT(
false, "Cannot convert ScalarType ", scalar_type, " to cudaDataType.")
}
}
} // namespace at::cuda
@@ -0,0 +1,204 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#pragma once
#if defined(PADDLE_WITH_HIP)
#include <hip/hip_runtime.h>
#elif defined(PADDLE_WITH_CUDA)
#include <cuda_runtime_api.h>
#endif
#include <c10/core/Device.h>
#include <c10/cuda/CUDAException.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <c10/util/Exception.h>
#include <memory>
#include <optional>
namespace at::cuda {
/**
* CUDAEvent is a movable, non-copyable wrapper around CUDA events.
* Provides compatibility with PyTorch's CUDAEvent API.
*/
struct CUDAEvent {
CUDAEvent() noexcept = default;
explicit CUDAEvent(unsigned int flags) noexcept : flags_(flags) {}
~CUDAEvent() {
if (is_created_) {
#ifdef PADDLE_WITH_HIP
hipEventDestroy(event_);
#else
cudaEventDestroy(event_);
#endif
}
}
CUDAEvent(const CUDAEvent&) = delete;
CUDAEvent& operator=(const CUDAEvent&) = delete;
CUDAEvent(CUDAEvent&& other) noexcept { moveHelper(std::move(other)); }
CUDAEvent& operator=(CUDAEvent&& other) noexcept {
if (this != &other) {
moveHelper(std::move(other));
}
return *this;
}
#ifdef PADDLE_WITH_HIP
operator hipEvent_t() const { return event(); }
hipEvent_t event() const { return event_; }
#else
operator cudaEvent_t() const { return event(); }
cudaEvent_t event() const { return event_; }
#endif
bool isCreated() const { return is_created_; }
c10::DeviceIndex device_index() const { return device_index_; }
bool query() const {
if (!is_created_) return true;
#ifdef PADDLE_WITH_HIP
hipError_t err = hipEventQuery(event_);
if (err == hipSuccess) return true;
if (err != hipErrorNotReady) C10_CUDA_CHECK(err);
#else
cudaError_t err = cudaEventQuery(event_);
if (err == cudaSuccess) return true;
if (err != cudaErrorNotReady) C10_CUDA_CHECK(err);
#endif
return false;
}
void record() { record(getCurrentCUDAStream()); }
void record(const CUDAStream& stream) {
if (!is_created_) {
createEvent(stream.unwrap().device_index());
}
TORCH_CHECK(device_index_ == stream.unwrap().device_index(),
"Event device ",
device_index_,
" does not match recording stream's device ",
stream.unwrap().device_index(),
".");
c10::cuda::CUDAGuard guard(device_index_);
#ifdef PADDLE_WITH_HIP
C10_CUDA_CHECK(hipEventRecord(event_, stream.stream()));
#else
C10_CUDA_CHECK(cudaEventRecord(event_, stream.stream()));
#endif
}
void recordOnce(const CUDAStream& stream) {
if (!was_recorded_) {
record(stream);
was_recorded_ = true;
}
}
void block(const CUDAStream& stream) {
if (is_created_) {
c10::cuda::CUDAGuard guard(stream.unwrap().device_index());
#ifdef PADDLE_WITH_HIP
C10_CUDA_CHECK(hipStreamWaitEvent(stream.stream(), event_, 0));
#else
C10_CUDA_CHECK(cudaStreamWaitEvent(stream.stream(), event_, 0));
#endif
}
}
void synchronize() const {
if (is_created_) {
#ifdef PADDLE_WITH_HIP
C10_CUDA_CHECK(hipEventSynchronize(event_));
#else
C10_CUDA_CHECK(cudaEventSynchronize(event_));
#endif
}
}
float elapsed_time(const CUDAEvent& other) const {
TORCH_CHECK(
is_created_ && other.isCreated(),
"Both events must be recorded before calculating elapsed time.");
TORCH_CHECK(
query() && other.query(),
"Both events must be completed before calculating elapsed time.");
float time_ms = 0;
c10::cuda::CUDAGuard guard(device_index_);
#ifdef PADDLE_WITH_HIP
C10_CUDA_CHECK(hipEventElapsedTime(&time_ms, event_, other.event_));
#else
C10_CUDA_CHECK(cudaEventElapsedTime(&time_ms, event_, other.event_));
#endif
return time_ms;
}
private:
#ifdef PADDLE_WITH_HIP
unsigned int flags_ = hipEventDisableTiming;
#else
unsigned int flags_ = cudaEventDisableTiming;
#endif
bool is_created_ = false;
bool was_recorded_ = false;
c10::DeviceIndex device_index_ = -1;
#ifdef PADDLE_WITH_HIP
hipEvent_t event_{};
#else
cudaEvent_t event_{};
#endif
void createEvent(c10::DeviceIndex device_index) {
device_index_ = device_index;
c10::cuda::CUDAGuard guard(device_index_);
#ifdef PADDLE_WITH_HIP
C10_CUDA_CHECK(hipEventCreateWithFlags(&event_, flags_));
#else
C10_CUDA_CHECK(cudaEventCreateWithFlags(&event_, flags_));
#endif
is_created_ = true;
}
void moveHelper(CUDAEvent&& other) {
flags_ = other.flags_;
is_created_ = std::exchange(other.is_created_, false);
was_recorded_ = other.was_recorded_;
device_index_ = other.device_index_;
#ifdef PADDLE_WITH_HIP
event_ = std::exchange(other.event_, hipEvent_t{});
#else
event_ = std::exchange(other.event_, cudaEvent_t{});
#endif
}
};
} // namespace at::cuda
namespace torch {
using at::cuda::CUDAEvent;
using at::cuda::CUDAStream;
} // namespace torch
@@ -0,0 +1,156 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#pragma once
#include <ATen/core/Generator.h>
#include <ATen/cuda/PhiloxCudaState.h>
#include <cstdint>
#include <memory>
#include <mutex>
#include <utility>
#include <vector>
#include "paddle/phi/backends/gpu/gpu_info.h"
#include "paddle/phi/core/generator.h"
namespace at {
// Forward declaration
struct CUDAGeneratorImpl;
inline DeviceIndex resolve_device_index(DeviceIndex idx) {
if (idx < 0) {
return static_cast<DeviceIndex>(phi::backends::gpu::GetCurrentDeviceId());
}
return idx;
}
struct CUDAGeneratorImpl : public c10::GeneratorImpl {
explicit CUDAGeneratorImpl(
DeviceIndex device_index = -1, // NOLINT(runtime/int)
bool use_default_gen = true)
: c10::GeneratorImpl(
c10::Device(c10::kCUDA, resolve_device_index(device_index)),
use_default_gen
? get_default_paddle_gen(resolve_device_index(device_index))
: create_new_paddle_gen(resolve_device_index(device_index))),
philox_offset_per_thread_(0) {}
~CUDAGeneratorImpl() override = default;
void set_current_seed(uint64_t seed) override {
gen_->SetCurrentSeed(seed);
philox_offset_per_thread_ = 0;
}
uint64_t current_seed() const override { return gen_->GetCurrentSeed(); }
uint64_t seed() override {
auto s = gen_->Seed();
philox_offset_per_thread_ = 0;
return s;
}
void set_offset(uint64_t offset) override {
philox_offset_per_thread_ = offset;
}
uint64_t get_offset() const override { return philox_offset_per_thread_; }
void set_philox_offset_per_thread(uint64_t offset) {
philox_offset_per_thread_ = offset;
}
uint64_t philox_offset_per_thread() const {
return philox_offset_per_thread_;
}
PhiloxCudaState philox_cuda_state(uint64_t increment) {
PhiloxCudaState state(gen_->GetCurrentSeed(), philox_offset_per_thread_);
philox_offset_per_thread_ += increment;
return state;
}
std::pair<uint64_t, uint64_t> philox_engine_inputs(uint64_t increment) {
uint64_t offset = philox_offset_per_thread_;
philox_offset_per_thread_ += increment;
return {gen_->GetCurrentSeed(), offset};
}
c10::intrusive_ptr<c10::GeneratorImpl> clone() const override {
auto new_gen = std::make_shared<phi::Generator>(gen_->GetCurrentSeed());
auto state = gen_->GetState();
new_gen->SetState(state);
auto impl = c10::make_intrusive<CUDAGeneratorImpl>(
static_cast<DeviceIndex>(device_.index()));
impl->gen_ = new_gen;
impl->philox_offset_per_thread_ = philox_offset_per_thread_;
return impl;
}
static c10::DeviceType device_type() { return c10::kCUDA; }
private:
uint64_t philox_offset_per_thread_;
static std::shared_ptr<phi::Generator> get_default_paddle_gen(
DeviceIndex device_index) {
return phi::DefaultCUDAGenerator(static_cast<int64_t>(device_index));
}
static std::shared_ptr<phi::Generator> create_new_paddle_gen(
DeviceIndex /*device_index*/) {
return std::make_shared<phi::Generator>();
}
};
namespace cuda {
namespace detail {
inline const Generator& getDefaultCUDAGenerator(DeviceIndex device_index = -1) {
auto idx = resolve_device_index(device_index);
static std::vector<Generator> generators;
static std::once_flag init_flag;
static int64_t num_devices = 0;
std::call_once(init_flag, []() {
num_devices = phi::backends::gpu::GetGPUDeviceCount();
generators.reserve(num_devices);
for (int64_t i = 0; i < num_devices; ++i) {
generators.emplace_back(c10::make_intrusive<CUDAGeneratorImpl>(
static_cast<DeviceIndex>(i), /*use_default_gen=*/true));
}
});
TORCH_CHECK(idx < static_cast<DeviceIndex>(num_devices),
"CUDA device index out of range: ",
idx);
return generators[static_cast<size_t>(idx)];
}
inline Generator createCUDAGenerator(DeviceIndex device_index = -1) {
return Generator(c10::make_intrusive<CUDAGeneratorImpl>(
device_index, /*use_default_gen=*/false));
}
} // namespace detail
} // namespace cuda
} // namespace at
@@ -0,0 +1,46 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <ATen/cuda/EmptyTensor.h>
#include "paddle/phi/api/include/api.h"
#include "paddle/phi/common/place.h"
namespace at::detail {
at::Tensor empty_cuda(IntArrayRef size,
ScalarType dtype,
std::optional<Device> device_opt,
std::optional<c10::MemoryFormat> memory_format_opt) {
PD_CHECK(!(memory_format_opt.has_value() &&
memory_format_opt.value() != c10::MemoryFormat::Contiguous),
"`MemoryFormat` other than Contiguous is not supported now.");
return paddle::experimental::empty(
size._PD_ToPaddleIntArray(),
compat::_PD_AtenScalarTypeToPhiDataType(dtype),
device_opt && device_opt->has_index() ? phi::GPUPlace(device_opt->index())
: paddle::DefaultGPUPlace());
}
at::Tensor empty_cuda(IntArrayRef size, const TensorOptions &options) {
auto place = options.has_device() && options.device().has_index()
? phi::GPUPlace(options.device().index())
: paddle::DefaultGPUPlace();
return paddle::experimental::empty(
size._PD_ToPaddleIntArray(),
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype_opt().value()),
place);
}
} // namespace at::detail
@@ -0,0 +1,32 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/TensorBody.h>
#include "paddle/common/macros.h"
namespace at::detail {
using at::Tensor;
PADDLE_API at::Tensor empty_cuda(
IntArrayRef size,
ScalarType dtype,
std::optional<Device> device_opt,
std::optional<c10::MemoryFormat> memory_format_opt);
PADDLE_API at::Tensor empty_cuda(IntArrayRef size,
const TensorOptions &options);
} // namespace at::detail
@@ -0,0 +1,16 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <c10/util/Exception.h>
@@ -0,0 +1,56 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#pragma once
#include <cstdint>
namespace at {
struct PhiloxCudaState {
PhiloxCudaState() = default;
// Called if graph capture is not underway
PhiloxCudaState(uint64_t seed, uint64_t offset) {
seed_.val = seed;
offset_.val = offset;
}
// Called if graph capture is underway
PhiloxCudaState(int64_t* seed,
int64_t* offset_extragraph,
uint64_t offset_intragraph) {
seed_.ptr = seed;
offset_.ptr = offset_extragraph;
offset_intragraph_ = offset_intragraph;
captured_ = true;
}
// Public members, directly accessible by at::cuda::philox::unpack.
// If we made them private with getters/setters, the getters/setters
// would have to be __device__, and we can't declare __device__ in ATen.
union Payload {
uint64_t val;
int64_t* ptr;
};
Payload seed_{};
Payload offset_{};
uint64_t offset_intragraph_ = 0;
bool captured_ = false;
};
} // namespace at
@@ -0,0 +1,49 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#pragma once
#include <ATen/cuda/PhiloxCudaState.h>
#include <tuple>
namespace at::cuda::philox {
// In-kernel call to retrieve philox seed and offset from a PhiloxCudaState
// instance whether that instance was created with graph capture underway or
// not. See Note [CUDA Graph-safe RNG states].
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
__host__ __device__ __forceinline__ std::tuple<uint64_t, uint64_t> unpack(
at::PhiloxCudaState arg) {
#else
inline std::tuple<uint64_t, uint64_t> unpack(at::PhiloxCudaState arg) {
#endif
if (arg.captured_) {
// static_cast avoids "warning: invalid narrowing conversion from "long" to
// "unsigned long".
// *(arg.offset_.ptr) is a broadcast load of a single int64_t to the entire
// kernel. For most threads' reads it will hit in cache, so it shouldn't
// hurt performance.
return std::make_tuple(
static_cast<uint64_t>(*arg.seed_.ptr),
static_cast<uint64_t>(*(arg.offset_.ptr) + arg.offset_intragraph_));
} else {
return std::make_tuple(arg.seed_.val, arg.offset_.val);
}
}
} // namespace at::cuda::philox
@@ -0,0 +1,75 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/AccumulateType.h>
#include <c10/core/Scalar.h>
#include <limits>
namespace at::native {
inline void arange_check_bounds(const c10::Scalar& start,
const c10::Scalar& end,
const c10::Scalar& step) {
// use double precision for validation to avoid precision issues
double dstart = start.to<double>();
double dend = end.to<double>();
double dstep = step.to<double>();
TORCH_CHECK(dstep > 0 || dstep < 0, "step must be nonzero");
TORCH_CHECK(std::isfinite(dstart) && std::isfinite(dend),
"unsupported range: ",
dstart,
" -> ",
dend);
TORCH_CHECK(
((dstep > 0) && (dend >= dstart)) || ((dstep < 0) && (dend <= dstart)),
"upper bound and lower bound inconsistent with step sign");
}
template <typename scalar_t>
int64_t compute_arange_size(const Scalar& start,
const Scalar& end,
const Scalar& step) {
arange_check_bounds(start, end, step);
// we use double precision for (start - end) / step
// to compute size_d for consistency across devices.
// The problem with using accscalar_t is that accscalar_t might be float32 on
// gpu for a float32 scalar_t, but double on cpu for the same, and the
// effective output size starts differing on CPU vs GPU because of precision
// issues, which we dont want. the corner-case we do want to take into account
// is int64_t, which has higher precision than double
double size_d;
if constexpr (std::is_same_v<scalar_t, int64_t>) {
using accscalar_t = at::acc_type<scalar_t, false>;
auto xstart = start.to<accscalar_t>();
auto xend = end.to<accscalar_t>();
auto xstep = step.to<accscalar_t>();
int64_t sgn = (xstep > 0) - (xstep < 0);
size_d = std::ceil((xend - xstart + xstep - sgn) / xstep);
} else {
size_d =
std::ceil((end.to<double>() - start.to<double>()) / step.to<double>());
}
TORCH_CHECK(size_d >= 0 && size_d <= static_cast<double>(
std::numeric_limits<int64_t>::max()),
"invalid size, possible overflow?");
return static_cast<int64_t>(size_d);
}
} // namespace at::native
@@ -0,0 +1,19 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
#include <c10/cuda/CUDAGuard.h>
#endif
@@ -0,0 +1,77 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/core/Scalar.h>
#include "paddle/phi/api/include/tensor.h"
#include "paddle/phi/common/data_type.h"
#include "paddle/phi/common/place.h"
namespace at {
/// Extracts a scalar value from a single-element dense tensor.
/// Mirrors PyTorch's at::_local_scalar_dense: copies the tensor to CPU if
/// needed, then reads the first element according to its dtype.
inline at::Scalar _local_scalar_dense(const at::Tensor& self) {
PD_CHECK(self.numel() > 0, "_local_scalar_dense: Empty tensor not supported");
// Move to CPU if necessary (for compatibility with PyTorch behavior)
const PaddleTensor& inner = self._PD_GetInner();
PaddleTensor cpu_tensor = inner;
if (!phi::is_cpu_place(inner.place())) {
PaddlePlace place(phi::AllocationType::CPU);
cpu_tensor = inner.copy_to(place, /*blocking=*/true);
}
auto dtype = cpu_tensor.dtype();
switch (dtype) {
case phi::DataType::FLOAT32:
return at::Scalar(*(cpu_tensor.data<float>()));
case phi::DataType::FLOAT64:
return at::Scalar(*(cpu_tensor.data<double>()));
case phi::DataType::FLOAT16:
return at::Scalar(
static_cast<float>(*(cpu_tensor.data<phi::dtype::float16>())));
case phi::DataType::BFLOAT16:
return at::Scalar(
static_cast<float>(*(cpu_tensor.data<phi::dtype::bfloat16>())));
case phi::DataType::INT8:
return at::Scalar(*(cpu_tensor.data<int8_t>()));
case phi::DataType::INT16:
return at::Scalar(*(cpu_tensor.data<int16_t>()));
case phi::DataType::INT32:
return at::Scalar(*(cpu_tensor.data<int32_t>()));
case phi::DataType::INT64:
return at::Scalar(*(cpu_tensor.data<int64_t>()));
case phi::DataType::UINT8:
return at::Scalar(*(cpu_tensor.data<uint8_t>()));
case phi::DataType::BOOL:
return at::Scalar(*(cpu_tensor.data<bool>()));
case phi::DataType::COMPLEX64:
return at::Scalar(*(cpu_tensor.data<phi::dtype::complex<float>>()));
case phi::DataType::COMPLEX128:
return at::Scalar(*(cpu_tensor.data<phi::dtype::complex<double>>()));
default:
PD_THROW("_local_scalar_dense: Unsupported data type");
}
}
} // namespace at
@@ -0,0 +1,44 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include "paddle/phi/core/sparse_coo_tensor.h"
#include "paddle/phi/core/sparse_csr_tensor.h"
namespace at {} // namespace at
namespace at {
inline int64_t Tensor::_nnz() const {
PD_CHECK(this->is_sparse(),
"_nnz expected sparse tensor layout but got ",
layout());
if (tensor_.layout() == common::DataLayout::SPARSE_COO) {
auto sparse_coo_tensor =
std::dynamic_pointer_cast<phi::SparseCooTensor>(tensor_.impl());
PD_CHECK(sparse_coo_tensor != nullptr,
"_nnz: failed to cast tensor impl to SparseCooTensor");
return sparse_coo_tensor->nnz();
} else {
auto sparse_csr_tensor =
std::dynamic_pointer_cast<phi::SparseCsrTensor>(tensor_.impl());
PD_CHECK(sparse_csr_tensor != nullptr,
"_nnz: failed to cast tensor impl to SparseCsrTensor");
return sparse_csr_tensor->nnz();
}
}
} // namespace at
@@ -0,0 +1,41 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include "paddle/phi/core/sparse_coo_tensor.h"
#include "paddle/phi/core/sparse_csr_tensor.h"
namespace at {} // namespace at
namespace at {
inline at::Tensor Tensor::_values() const {
PD_CHECK(this->is_sparse(),
"_values expected sparse tensor layout but got ",
layout());
if (tensor_.layout() == common::DataLayout::SPARSE_COO) {
auto sparse_coo_tensor =
std::dynamic_pointer_cast<phi::SparseCooTensor>(tensor_.impl());
PD_CHECK(sparse_coo_tensor != nullptr,
"_values: failed to cast tensor impl to SparseCooTensor");
return paddle::Tensor(
std::make_shared<phi::DenseTensor>(sparse_coo_tensor->values()));
} else {
PD_THROW("_values is not implemented for SparseCsr tensors");
}
}
} // namespace at
@@ -0,0 +1,57 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/core/TensorOptions.h>
#include <optional>
#include <string_view>
#include "paddle/phi/api/include/api.h"
#include "paddle/phi/core/enforce.h"
namespace at {
inline at::Tensor abs(const at::Tensor& self) {
if (!self.is_contiguous()) {
phi::enforce::ThrowWarnInternal(
"at::abs: input tensor is non-contiguous. PyTorch and Paddle handle "
"non-contiguous tensors differently, which may produce logically "
"incorrect results even though the code is syntactically valid. "
"See https://github.com/PaddlePaddle/Paddle/pull/78099 for details.");
}
return paddle::experimental::abs(self._PD_GetInner());
}
} // namespace at
namespace at {
inline at::Tensor Tensor::abs() const { return at::abs(*this); }
inline at::Tensor& Tensor::abs_() const {
if (!is_contiguous()) {
phi::enforce::ThrowWarnInternal(
"Tensor::abs_: tensor is non-contiguous. PyTorch and Paddle handle "
"non-contiguous tensors differently, which may produce logically "
"incorrect results even though the code is syntactically valid. "
"See https://github.com/PaddlePaddle/Paddle/pull/78099 for details.");
}
PaddleTensor& inner = const_cast<PaddleTensor&>(tensor_);
paddle::experimental::abs_(inner);
return const_cast<at::Tensor&>(*this);
}
} // namespace at
@@ -0,0 +1,63 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/util/OptionalArrayRef.h>
#include "paddle/phi/api/include/api.h"
namespace at {
// all: Check if all elements are true (non-zero)
// Version 1: all() - check all elements in the tensor
inline at::Tensor all(const at::Tensor& self) {
return paddle::experimental::all(self._PD_GetInner(), {}, false);
}
// Version 2: all(dim, keepdim) - check along a specific dimension
inline at::Tensor all(const at::Tensor& self,
int64_t dim,
bool keepdim = false) {
return paddle::experimental::all(self._PD_GetInner(), {dim}, keepdim);
}
// Version 3: all(dim, keepdim) - check along optional dimensions
inline at::Tensor all(const at::Tensor& self,
at::OptionalIntArrayRef dim,
bool keepdim = false) {
std::vector<int64_t> axis_vec;
if (dim.has_value()) {
axis_vec.assign(dim.value().begin(), dim.value().end());
}
return paddle::experimental::all(self._PD_GetInner(), axis_vec, keepdim);
}
} // namespace at
namespace at {
// Tensor member function implementations
inline at::Tensor Tensor::all() const { return at::all(*this); }
inline at::Tensor Tensor::all(int64_t dim, bool keepdim) const {
return at::all(*this, dim, keepdim);
}
inline at::Tensor Tensor::all(at::OptionalIntArrayRef dim, bool keepdim) const {
return at::all(*this, dim, keepdim);
}
} // namespace at
@@ -0,0 +1,57 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include "paddle/phi/api/include/api.h"
#include "paddle/phi/common/data_type.h"
#include "paddle/phi/common/scalar.h"
namespace at {
// allclose: Check if two tensors are close to each other
inline bool allclose(const at::Tensor& self,
const at::Tensor& other,
double rtol = 1e-05,
double atol = 1e-08,
bool equal_nan = false) {
// Paddle's allclose returns a Tensor, but PyTorch's allclose returns bool.
// The allclose kernel always sets output dtype to phi::DataType::BOOL via
// kernel->OutputAt(0).SetDataType(phi::DataType::BOOL), so we read BOOL
// directly.
PaddleTensor result = paddle::experimental::allclose(self._PD_GetInner(),
other._PD_GetInner(),
phi::Scalar(rtol),
phi::Scalar(atol),
equal_nan);
auto* result_tensor =
std::dynamic_pointer_cast<phi::DenseTensor>(result.impl()).get();
return *result_tensor->data<bool>();
}
} // namespace at
namespace at {
// Tensor member function implementation
inline bool Tensor::allclose(const at::Tensor& other,
double rtol,
double atol,
bool equal_nan) const {
return at::allclose(*this, other, rtol, atol, equal_nan);
}
} // namespace at
@@ -0,0 +1,61 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/core/Scalar.h>
#include <c10/util/OptionalArrayRef.h>
#include "paddle/phi/api/include/api.h"
namespace at {
// any - free functions
inline Tensor any(const Tensor& self, int64_t dim, bool keepdim = false) {
return paddle::experimental::any(self._PD_GetInner(), {dim}, keepdim);
}
inline Tensor any(const Tensor& self,
at::OptionalIntArrayRef dim,
bool keepdim = false) {
std::vector<int64_t> dims_vec;
if (dim.has_value() && dim.value().size() > 0) {
dims_vec.assign(dim.value().begin(), dim.value().end());
}
return paddle::experimental::any(self._PD_GetInner(), dims_vec, keepdim);
}
inline Tensor any(const Tensor& self) {
return paddle::experimental::any(self._PD_GetInner());
}
// any - member function implementations
inline Tensor Tensor::any(int64_t dim, bool keepdim) const {
return paddle::experimental::any(_PD_GetInner(), {dim}, keepdim);
}
inline Tensor Tensor::any(at::OptionalIntArrayRef dim, bool keepdim) const {
std::vector<int64_t> dims_vec;
if (dim.has_value() && dim.value().size() > 0) {
dims_vec.assign(dim.value().begin(), dim.value().end());
}
return paddle::experimental::any(_PD_GetInner(), dims_vec, keepdim);
}
inline Tensor Tensor::any() const {
return paddle::experimental::any(_PD_GetInner());
}
} // namespace at
@@ -0,0 +1,157 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#pragma once
#include <ATen/core/Tensor.h>
#include <ATen/native/RangeUtils.h>
#include <c10/core/TensorOptions.h>
#include <utils/pinned_place.h>
#include <optional>
#include "paddle/phi/api/include/api.h"
#include "paddle/phi/common/place.h"
namespace at {
namespace detail {
inline bool _PD_IsIntegralArangeScalar(const at::Scalar& scalar) {
switch (scalar.dtype()) {
case phi::DataType::BOOL:
case phi::DataType::UINT8:
case phi::DataType::INT8:
case phi::DataType::UINT16:
case phi::DataType::INT16:
case phi::DataType::UINT32:
case phi::DataType::INT32:
case phi::DataType::UINT64:
case phi::DataType::INT64:
return true;
default:
return false;
}
}
inline at::ScalarType _PD_ResolveArangeDtype(const at::Scalar& start,
const at::Scalar& end,
const at::Scalar& step,
const at::TensorOptions& options) {
if (options.has_dtype()) {
return options.dtype().toScalarType();
}
if (_PD_IsIntegralArangeScalar(start) && _PD_IsIntegralArangeScalar(end) &&
_PD_IsIntegralArangeScalar(step)) {
return at::kLong;
}
return c10::get_default_dtype_as_scalartype();
}
inline paddle::Tensor _PD_MakeArangeScalarTensor(const at::Scalar& scalar,
phi::DataType dtype) {
return paddle::experimental::full({}, scalar, dtype, phi::CPUPlace());
}
} // namespace detail
inline at::Tensor arange(const at::Scalar& start,
const at::Scalar& end,
const at::Scalar& step,
at::TensorOptions options = {}) {
// Match PyTorch: step must be non-zero and consistent with (end - start).
at::native::arange_check_bounds(start, end, step);
auto dtype = detail::_PD_ResolveArangeDtype(start, end, step, options);
auto pd_dtype = compat::_PD_AtenScalarTypeToPhiDataType(dtype);
if (options.pinned_memory()) {
// Pinning memory is only supported for CPU tensors
if (options.has_device() && !options.device().is_cpu()) {
PD_THROW(
"pin_memory=true requires device to be CPU, but got non-CPU device");
}
phi::Place base_place = options._PD_GetPlace();
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(base_place);
auto dense = paddle::experimental::arange(
detail::_PD_MakeArangeScalarTensor(start, pd_dtype),
detail::_PD_MakeArangeScalarTensor(end, pd_dtype),
detail::_PD_MakeArangeScalarTensor(step, pd_dtype),
pd_dtype,
phi::CPUPlace());
return dense.copy_to(pinned_place, /*blocking=*/true);
}
return paddle::experimental::arange(
detail::_PD_MakeArangeScalarTensor(start, pd_dtype),
detail::_PD_MakeArangeScalarTensor(end, pd_dtype),
detail::_PD_MakeArangeScalarTensor(step, pd_dtype),
pd_dtype,
options._PD_GetPlace());
}
inline at::Tensor arange(const at::Scalar& end,
at::TensorOptions options = {}) {
return arange(/*start=*/0, end, /*step=*/1, options);
}
inline at::Tensor arange(const at::Scalar& end,
::std::optional<at::ScalarType> dtype,
::std::optional<at::Layout> layout,
::std::optional<at::Device> device,
::std::optional<bool> pin_memory) {
auto options = at::TensorOptions()
.dtype(dtype)
.layout(layout)
.device(device)
.pinned_memory(pin_memory);
return arange(/*start=*/0, end, /*step=*/1, options);
}
inline at::Tensor arange(const at::Scalar& start,
const at::Scalar& end,
at::TensorOptions options = {}) {
return arange(start, end, /*step=*/1, options);
}
inline at::Tensor arange(const at::Scalar& start,
const at::Scalar& end,
::std::optional<at::ScalarType> dtype,
::std::optional<at::Layout> layout,
::std::optional<at::Device> device,
::std::optional<bool> pin_memory) {
auto options = at::TensorOptions()
.dtype(dtype)
.layout(layout)
.device(device)
.pinned_memory(pin_memory);
return arange(start, end, /*step=*/1, options);
}
inline at::Tensor arange(const at::Scalar& start,
const at::Scalar& end,
const at::Scalar& step,
::std::optional<at::ScalarType> dtype,
::std::optional<at::Layout> layout,
::std::optional<at::Device> device,
::std::optional<bool> pin_memory) {
auto options = at::TensorOptions()
.dtype(dtype)
.layout(layout)
.device(device)
.pinned_memory(pin_memory);
return arange(start, end, step, options);
}
} // namespace at
@@ -0,0 +1,113 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/util/ArrayRef.h>
#include <optional>
#include <vector>
#include "paddle/common/ddim.h"
#include "paddle/phi/core/dense_tensor.h"
namespace at {
// as_strided: Create a tensor view with custom size, stride, and storage_offset
inline at::Tensor Tensor::as_strided(
at::IntArrayRef size,
at::IntArrayRef stride,
::std::optional<int64_t> storage_offset) const {
// Materialize the compat StorageHolderView before creating the view so
// aliasing tensors share one StorageImpl and observe later resize_ growth.
(void)this->storage();
auto src_impl = tensor_.impl();
auto* src_tensor =
std::dynamic_pointer_cast<phi::DenseTensor>(src_impl).get();
if (!src_tensor) {
PD_THROW("as_strided: tensor must be a DenseTensor");
}
// Create new meta with desired shape and strides first
std::vector<int64_t> size_vec(size.begin(), size.end());
std::vector<int64_t> stride_vec(stride.begin(), stride.end());
// Create new DenseTensor with correct meta, then share data
// We need to create a temporary DenseTensor with the right meta
// because ShareDataWith copies the source meta which we don't want
auto new_tensor = std::make_shared<phi::DenseTensor>();
// First, set up the holder by sharing data (this copies src meta, we'll
// override)
new_tensor->ShareDataWith(*src_tensor);
// Now create the correct meta with new shape/strides
phi::DenseTensorMeta meta(src_tensor->dtype(),
common::make_ddim(size_vec),
common::make_ddim(stride_vec));
// Calculate offset in bytes
int64_t offset = storage_offset.has_value() ? storage_offset.value() : 0;
meta.offset = src_tensor->meta().offset +
static_cast<size_t>(offset) * phi::SizeOf(src_tensor->dtype());
new_tensor->set_meta(meta);
PaddleTensor result;
result.set_impl(new_tensor);
return Tensor(result);
}
// as_strided_: Inplace version
inline const at::Tensor& Tensor::as_strided_(
at::IntArrayRef size,
at::IntArrayRef stride,
::std::optional<int64_t> storage_offset) const {
// Keep inplace metadata-only view rewrites attached to the same compat
// storage as the original tensor.
(void)this->storage();
auto src_impl = tensor_.impl();
auto* src_tensor =
std::dynamic_pointer_cast<phi::DenseTensor>(src_impl).get();
if (!src_tensor) {
PD_THROW("as_strided_: tensor must be a DenseTensor");
}
std::vector<int64_t> size_vec(size.begin(), size.end());
std::vector<int64_t> stride_vec(stride.begin(), stride.end());
// Use set_meta instead of Resize + set_strides to avoid contiguous check
phi::DenseTensorMeta meta(src_tensor->dtype(),
common::make_ddim(size_vec),
common::make_ddim(stride_vec));
meta.layout = src_tensor->layout();
int64_t offset = storage_offset.has_value() ? storage_offset.value() : 0;
meta.offset = src_tensor->meta().offset +
static_cast<size_t>(offset) * phi::SizeOf(src_tensor->dtype());
src_tensor->set_meta(meta);
return *this;
}
// as_strided_scatter: Scatter src into a strided view
// Returns a new tensor (copy of self) with the strided window filled by src.
// The original tensor is NOT modified.
inline at::Tensor Tensor::as_strided_scatter(
const at::Tensor& src,
at::IntArrayRef size,
at::IntArrayRef stride,
::std::optional<int64_t> storage_offset) const {
// Clone self to an independent copy so the original tensor is left unchanged
PaddleTensor self_copy = tensor_.copy_to(tensor_.place(), /*blocking=*/true);
at::Tensor copy_tensor(self_copy);
at::Tensor strided_view =
copy_tensor.as_strided(size, stride, storage_offset);
strided_view.copy_(src);
return copy_tensor;
}
} // namespace at
@@ -0,0 +1,34 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/core/TensorOptions.h>
#include <optional>
#include "paddle/phi/api/include/api.h"
namespace at {
inline at::Tensor cat(const at::ITensorListRef& tensors, int64_t dim = 0) {
std::vector<paddle::Tensor> pd_tensors;
pd_tensors.reserve(tensors.size());
for (const auto& t : tensors) {
pd_tensors.push_back(t._PD_GetInner());
}
return paddle::experimental::concat(pd_tensors, dim);
}
} // namespace at
@@ -0,0 +1,98 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <vector>
#include "paddle/phi/api/include/api.h"
namespace at {
// chunk - splits tensor into chunks
inline std::vector<Tensor> chunk(const Tensor& self,
int64_t chunks,
int64_t dim = 0) {
if (chunks <= 0) {
PD_THROW("chunk expects chunks to be greater than 0, got ", chunks);
}
std::vector<Tensor> result;
paddle::Tensor pd_tensor = self._PD_GetInner();
int64_t rank = static_cast<int64_t>(pd_tensor.dims().size());
if (rank == 0) {
PD_THROW("chunk expects at least a 1-dimensional tensor");
}
int64_t original_dim = dim;
if (dim < 0) {
dim += rank;
}
if (dim < 0 || dim >= rank) {
PD_THROW("Dimension out of range (expected to be in range of [",
-rank,
", ",
rank - 1,
"], but got ",
original_dim,
")");
}
int64_t dim_size = pd_tensor.dims()[dim];
if (dim_size == 0) {
for (int64_t i = 0; i < chunks; ++i) {
auto chunk_tensor =
paddle::experimental::slice(pd_tensor, {dim}, {0}, {0}, {1}, {});
result.push_back(Tensor(chunk_tensor));
}
return result;
}
// PyTorch returns at most 'dim_size' non-empty chunks when chunks > dim_size
if (chunks > dim_size) {
chunks = dim_size;
}
int64_t chunk_size = (dim_size + chunks - 1) / chunks;
int64_t remaining = dim_size;
for (int64_t i = 0; i < chunks && remaining > 0; ++i) {
int64_t current_chunk_size = std::min(chunk_size, remaining);
auto chunk_tensor =
paddle::experimental::slice(pd_tensor,
{dim},
{i * chunk_size},
{i * chunk_size + current_chunk_size},
{1},
{});
result.push_back(Tensor(chunk_tensor));
remaining -= current_chunk_size;
}
return result;
}
} // namespace at
namespace at {
// Member function: Tensor::chunk
inline std::vector<Tensor> Tensor::chunk(int64_t chunks, int64_t dim) const {
return at::chunk(*this, chunks, dim);
}
} // namespace at
@@ -0,0 +1,212 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <ATen/core/TensorBase.h>
#include <ATen/ops/full.h>
#include <c10/core/Scalar.h>
#include <c10/core/ScalarType.h>
#include <limits>
#include <optional>
#include "paddle/phi/api/include/tensor.h"
namespace at {
// Helper function implementations
namespace detail {
inline at::Scalar get_default_min_value(c10::ScalarType dtype) {
switch (dtype) {
case c10::ScalarType::Byte:
return at::Scalar(static_cast<uint8_t>(0));
case c10::ScalarType::Char:
return at::Scalar(std::numeric_limits<int8_t>::lowest());
case c10::ScalarType::Short:
return at::Scalar(std::numeric_limits<int16_t>::lowest());
case c10::ScalarType::Int:
return at::Scalar(std::numeric_limits<int32_t>::lowest());
case c10::ScalarType::Long:
return at::Scalar(std::numeric_limits<int64_t>::lowest());
case c10::ScalarType::UInt16:
return at::Scalar(static_cast<uint16_t>(0));
case c10::ScalarType::UInt32:
return at::Scalar(static_cast<uint32_t>(0));
case c10::ScalarType::UInt64:
return at::Scalar(static_cast<uint64_t>(0));
case c10::ScalarType::Half:
return at::Scalar(-std::numeric_limits<float>::infinity());
case c10::ScalarType::Float:
return at::Scalar(-std::numeric_limits<float>::infinity());
case c10::ScalarType::Double:
return at::Scalar(-std::numeric_limits<double>::infinity());
case c10::ScalarType::BFloat16:
return at::Scalar(-std::numeric_limits<float>::infinity());
case c10::ScalarType::Bool:
return at::Scalar(false);
default:
return at::Scalar(-std::numeric_limits<double>::infinity());
}
}
inline at::Scalar get_default_max_value(c10::ScalarType dtype) {
switch (dtype) {
case c10::ScalarType::Byte:
return at::Scalar(std::numeric_limits<uint8_t>::max());
case c10::ScalarType::Char:
return at::Scalar(std::numeric_limits<int8_t>::max());
case c10::ScalarType::Short:
return at::Scalar(std::numeric_limits<int16_t>::max());
case c10::ScalarType::Int:
return at::Scalar(std::numeric_limits<int32_t>::max());
case c10::ScalarType::Long:
return at::Scalar(std::numeric_limits<int64_t>::max());
case c10::ScalarType::UInt16:
return at::Scalar(std::numeric_limits<uint16_t>::max());
case c10::ScalarType::UInt32:
return at::Scalar(std::numeric_limits<uint32_t>::max());
case c10::ScalarType::UInt64:
return at::Scalar(std::numeric_limits<uint64_t>::max());
case c10::ScalarType::Half:
return at::Scalar(std::numeric_limits<float>::infinity());
case c10::ScalarType::Float:
return at::Scalar(std::numeric_limits<float>::infinity());
case c10::ScalarType::Double:
return at::Scalar(std::numeric_limits<double>::infinity());
case c10::ScalarType::BFloat16:
return at::Scalar(std::numeric_limits<float>::infinity());
case c10::ScalarType::Bool:
return at::Scalar(true);
default:
return at::Scalar(std::numeric_limits<double>::infinity());
}
}
} // namespace detail
} // namespace at
namespace at {
inline at::Tensor Tensor::clamp(const ::std::optional<at::Scalar>& min,
const ::std::optional<at::Scalar>& max) const {
// Handle cases where min or max is nullopt - don't apply that bound
if (min.has_value() && !max.has_value()) {
// Only min is specified - use clamp_min
return clamp_min(min.value());
} else if (!min.has_value() && max.has_value()) {
// Only max is specified - use clamp_max
return clamp_max(max.value());
} else if (!min.has_value() && !max.has_value()) {
// Neither specified - return copy of tensor
return *this;
}
// Both specified - apply full clamp
return Tensor(paddle::experimental::clip(tensor_, min.value(), max.value()));
}
inline at::Tensor Tensor::clamp(const ::std::optional<at::Tensor>& min,
const ::std::optional<at::Tensor>& max) const {
PaddleTensor result = tensor_;
if (min.has_value()) {
result = paddle::experimental::maximum(result, min.value()._PD_GetInner());
}
if (max.has_value()) {
result = paddle::experimental::minimum(result, max.value()._PD_GetInner());
}
return Tensor(result);
}
inline at::Tensor& Tensor::clamp_(
const ::std::optional<at::Scalar>& min,
const ::std::optional<at::Scalar>& max) const {
// Handle cases where min or max is nullopt - don't apply that bound
if (min.has_value() && !max.has_value()) {
// Only min is specified - use clamp_min_
return clamp_min_(min.value());
} else if (!min.has_value() && max.has_value()) {
// Only max is specified - use clamp_max_
return clamp_max_(max.value());
} else if (!min.has_value() && !max.has_value()) {
// Neither specified - nothing to do
return const_cast<at::Tensor&>(*this);
}
// Both specified - apply full clamp
paddle::experimental::clip_(
const_cast<PaddleTensor&>(tensor_), min.value(), max.value());
return const_cast<at::Tensor&>(*this);
}
inline at::Tensor& Tensor::clamp_(
const ::std::optional<at::Tensor>& min,
const ::std::optional<at::Tensor>& max) const {
if (min.has_value()) {
PaddleTensor temp =
paddle::experimental::maximum(tensor_, min.value()._PD_GetInner());
const_cast<PaddleTensor&>(tensor_) = temp;
}
if (max.has_value()) {
PaddleTensor temp =
paddle::experimental::minimum(tensor_, max.value()._PD_GetInner());
const_cast<PaddleTensor&>(tensor_) = temp;
}
return const_cast<at::Tensor&>(*this);
}
inline at::Tensor Tensor::clamp_max(const at::Scalar& max) const {
// Create a tensor with the same shape filled with the max value
at::Tensor max_tensor = at::full(tensor_.shape(), max, {});
return clamp_max(max_tensor);
}
inline at::Tensor Tensor::clamp_max(const at::Tensor& max) const {
return Tensor(paddle::experimental::minimum(tensor_, max._PD_GetInner()));
}
inline at::Tensor& Tensor::clamp_max_(const at::Scalar& max) const {
// Create a tensor with the same shape filled with the max value
at::Tensor max_tensor = at::full(tensor_.shape(), max, {});
return clamp_max_(max_tensor);
}
inline at::Tensor& Tensor::clamp_max_(const at::Tensor& max) const {
PaddleTensor temp =
paddle::experimental::minimum(tensor_, max._PD_GetInner());
const_cast<PaddleTensor&>(tensor_) = temp;
return const_cast<at::Tensor&>(*this);
}
inline at::Tensor Tensor::clamp_min(const at::Scalar& min) const {
// Create a tensor with the same shape filled with the min value
at::Tensor min_tensor = at::full(tensor_.shape(), min, {});
return clamp_min(min_tensor);
}
inline at::Tensor Tensor::clamp_min(const at::Tensor& min) const {
return Tensor(paddle::experimental::maximum(tensor_, min._PD_GetInner()));
}
inline at::Tensor& Tensor::clamp_min_(const at::Scalar& min) const {
// Create a tensor with the same shape filled with the min value
at::Tensor min_tensor = at::full(tensor_.shape(), min, {});
return clamp_min_(min_tensor);
}
inline at::Tensor& Tensor::clamp_min_(const at::Tensor& min) const {
PaddleTensor temp =
paddle::experimental::maximum(tensor_, min._PD_GetInner());
const_cast<PaddleTensor&>(tensor_) = temp;
return const_cast<at::Tensor&>(*this);
}
} // namespace at
@@ -0,0 +1,34 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include "paddle/phi/api/include/sparse_api.h"
namespace at {} // namespace at
namespace at {
inline at::Tensor Tensor::coalesce() const {
PD_CHECK(layout() == kSparse,
"coalesce expected sparse coordinate tensor layout but got ",
layout());
if (is_coalesced()) {
return *this;
}
return paddle::experimental::sparse::coalesce(tensor_);
}
} // namespace at
@@ -0,0 +1,43 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
namespace at {
inline at::Tensor detach(const at::Tensor& self) {
// Create a new Tensor that shares data but has no autograd history
auto inner = self._PD_GetInner();
PaddleTensor detached_tensor(inner.impl());
detached_tensor.set_name(inner.name());
detached_tensor.set_autograd_meta(nullptr);
return Tensor(detached_tensor);
}
} // namespace at
namespace at {
inline at::Tensor Tensor::detach() const { return at::detach(*this); }
inline at::Tensor& Tensor::detach_() const {
// In-place version: clear autograd meta of current tensor
PaddleTensor& inner = const_cast<PaddleTensor&>(tensor_);
inner.set_autograd_meta(nullptr);
return const_cast<at::Tensor&>(*this);
}
} // namespace at
@@ -0,0 +1,44 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <ATen/ops/tensor_split.h>
namespace at {
inline std::vector<at::Tensor> dsplit(const at::Tensor& self,
int64_t sections) {
return tensor_split(self, sections, 2);
}
inline std::vector<at::Tensor> dsplit(const at::Tensor& self,
at::IntArrayRef indices) {
return tensor_split(self, indices, 2);
}
} // namespace at
namespace at {
inline std::vector<at::Tensor> Tensor::dsplit(int64_t sections) const {
return at::dsplit(*this, sections);
}
inline std::vector<at::Tensor> Tensor::dsplit(at::IntArrayRef indices) const {
return at::dsplit(*this, indices);
}
} // namespace at
@@ -0,0 +1,75 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/core/TensorOptions.h>
#include <utils/dense_sparse_conversion.h>
#include <optional>
#include <string_view>
#include "paddle/phi/api/include/api.h"
#include "paddle/phi/common/place.h"
namespace at {
inline at::Tensor empty(
at::IntArrayRef size,
at::TensorOptions options = {},
::std::optional<at::MemoryFormat> memory_format = ::std::nullopt) {
PD_CHECK(!(memory_format.has_value() &&
memory_format.value() != c10::MemoryFormat::Contiguous),
"`MemoryFormat` other than Contiguous is not supported now.");
if (options.pinned_memory()) {
// Pinning memory is only supported for CPU tensors
if (options.has_device() && !options.device().is_cpu()) {
PD_THROW(
"pin_memory=true requires device to be CPU, but got non-CPU device");
}
auto dense = paddle::experimental::empty(
size._PD_ToPaddleIntArray(),
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
phi::CPUPlace());
dense = dense.copy_to(phi::GPUPinnedPlace(), /*blocking=*/true);
return compat::_PD_ConvertToSparseIfNeeded(dense, options.layout());
}
auto dense = paddle::experimental::empty(
size._PD_ToPaddleIntArray(),
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
options._PD_GetPlace());
return compat::_PD_ConvertToSparseIfNeeded(dense, options.layout());
}
inline at::Tensor empty(at::IntArrayRef size,
::std::optional<at::ScalarType> dtype,
::std::optional<at::Layout> layout,
::std::optional<at::Device> device,
::std::optional<bool> pin_memory,
::std::optional<at::MemoryFormat> memory_format) {
PD_CHECK(!(memory_format.has_value() &&
memory_format.value() != c10::MemoryFormat::Contiguous),
"`MemoryFormat` other than Contiguous is not supported now.");
auto options =
at::TensorOptions()
.dtype(dtype.value_or(c10::get_default_dtype_as_scalartype()))
.layout(layout)
.device(device.value_or(at::kCPU))
.pinned_memory(pin_memory);
return empty(size, options, memory_format);
}
#define empty_symint empty // SymIntArrayRef is same as IntArrayRef
} // namespace at
@@ -0,0 +1,81 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/core/TensorOptions.h>
#include <utils/dense_sparse_conversion.h>
#include <utils/pinned_place.h>
#include <optional>
#include <string_view>
#include "paddle/phi/api/include/api.h"
#include "paddle/phi/common/place.h"
namespace at {
inline at::Tensor empty_like(
const at::Tensor& self,
at::TensorOptions options = {},
::std::optional<at::MemoryFormat> memory_format = ::std::nullopt) {
PD_CHECK(!(memory_format.has_value() &&
memory_format.value() != c10::MemoryFormat::Contiguous),
"`MemoryFormat` other than Contiguous is not supported now.");
auto dtype = options.dtype_opt().value_or(self.dtype());
paddle::Tensor dense;
if (options.pinned_memory()) {
// Pinning memory is only supported for CPU tensors
if (options.has_device() && !options.device().is_cpu()) {
PD_THROW(
"pin_memory=true requires device to be CPU, but got non-CPU device");
}
auto dense_cpu = paddle::experimental::empty_like(
self._PD_GetInner(),
compat::_PD_AtenScalarTypeToPhiDataType(dtype),
phi::CPUPlace());
phi::Place base_place = options._PD_GetPlace();
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(base_place);
dense = dense_cpu.copy_to(pinned_place, /*blocking=*/true);
} else {
auto place = options.device_opt().value_or(self.device());
dense = paddle::experimental::empty_like(
self._PD_GetInner(),
compat::_PD_AtenScalarTypeToPhiDataType(dtype),
place._PD_GetInner());
}
return compat::_PD_ConvertToSparseIfNeeded(dense, options.layout());
}
inline at::Tensor empty_like(const at::Tensor& self,
::std::optional<at::ScalarType> dtype,
::std::optional<at::Layout> layout,
::std::optional<at::Device> device,
::std::optional<bool> pin_memory,
::std::optional<at::MemoryFormat> memory_format) {
PD_CHECK(!(memory_format.has_value() &&
memory_format.value() != c10::MemoryFormat::Contiguous),
"`MemoryFormat` other than Contiguous is not supported now.");
auto options = at::TensorOptions()
.dtype(dtype)
.layout(layout)
.device(device)
.pinned_memory(pin_memory);
return empty_like(self, options, memory_format);
}
} // namespace at
@@ -0,0 +1,41 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/core/TensorOptions.h>
#include <c10/util/ArrayRef.h>
#include <optional>
#include <string_view>
#include "paddle/phi/api/include/api.h"
namespace at {
inline at::Tensor empty_strided(at::IntArrayRef size,
at::IntArrayRef stride,
at::TensorOptions options = {}) {
auto empty_tensor = paddle::experimental::empty(
size._PD_ToPaddleIntArray(),
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
options._PD_GetPlace());
return paddle::experimental::as_strided(
empty_tensor,
std::vector<int64_t>(size.begin(), size.end()),
std::vector<int64_t>(stride.begin(), stride.end()));
}
} // namespace at
@@ -0,0 +1,60 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <ATen/ops/item.h>
#include "paddle/phi/api/include/api.h"
namespace at {
inline bool equal(const at::Tensor& self, const at::Tensor& other) {
PD_CHECK(self.defined(),
"Expected a proper Tensor but got None (or an undefined Tensor in "
"C++)");
PD_CHECK(other.defined(),
"Expected a proper Tensor but got None (or an undefined Tensor in "
"C++)");
PD_CHECK(self.device() == other.device(),
"Cannot compare two tensors on "
"different devices. Got: ",
self.device(),
" and ",
other.device());
if (self.sizes() != other.sizes()) {
return false;
}
auto lhs = self._PD_GetInner();
auto rhs = other._PD_GetInner();
if (self.scalar_type() != other.scalar_type()) {
rhs = paddle::experimental::cast(
rhs, compat::_PD_AtenScalarTypeToPhiDataType(self.scalar_type()));
}
auto result = paddle::experimental::equal_all(lhs, rhs);
return at::Tensor(std::move(result)).item<bool>();
}
} // namespace at
namespace at {
inline bool Tensor::equal(const at::Tensor& other) const {
return at::equal(*this, other);
}
} // namespace at
@@ -0,0 +1,137 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include "paddle/phi/api/include/api.h"
namespace at {
// expand - expands tensor to new size
// PyTorch's expand works by right-aligning dimensions and broadcasting
// dimensions with size 1 to the target size
// Unlike Paddle's expand_v2, PyTorch allows non-singleton dimensions to be
// preserved when they match the corresponding target dimension
inline Tensor expand(const Tensor& self,
at::IntArrayRef size,
bool implicit = false) {
// implicit parameter is used by PyTorch's vmap for internal optimization.
// It doesn't affect the actual expand operation, so we can safely ignore it.
paddle::Tensor pd_tensor = self._PD_GetInner();
// Target sizes - convert to vector
std::vector<int64_t> target_size_vec(size.begin(), size.end());
auto target_rank = target_size_vec.size();
auto input_dims = pd_tensor.dims();
auto input_rank = static_cast<size_t>(input_dims.size());
// PyTorch's expand uses right-alignment semantics:
// - For 1D tensor expand to 2D: {3}.expand({3,4}) treats input as {3,1},
// expands to {3,4}
// - Non-singleton dimensions are preserved, singleton dimensions (1) can
// expand
//
// For example:
// {3}.expand({3, 4}) -> input {3} becomes {3, 1} implicitly
// then expand: dim 0: 3 stays 3, dim 1: 1 -> 4 -> result {3, 4}
if (input_rank < target_rank) {
// Add leading 1s to right-align with target shape (PyTorch behavior)
// Input {1, 2}, target {2, 3, 2} -> reshape to {1, 1, 2}
std::vector<int64_t> reshape_vec(target_rank, 1);
for (size_t i = 0; i < input_rank; ++i) {
reshape_vec[target_rank - input_rank + i] = input_dims[i];
}
// Check if Paddle's expand can handle this right-aligned shape
// Paddle allows: input[i] == 1 (can expand), or input[i] == target[i]
// (match)
bool can_use_paddle_expand = true;
size_t fail_dim = 0;
for (size_t i = 0; i < target_rank; ++i) {
bool dim_can_expand = (reshape_vec[i] == 1);
bool dim_is_matching = (reshape_vec[i] == target_size_vec[i]);
if (!dim_can_expand && !dim_is_matching) {
can_use_paddle_expand = false;
fail_dim = i;
break;
}
}
if (can_use_paddle_expand) {
// Reshape to right-aligned shape, then expand
paddle::Tensor reshaped =
paddle::experimental::reshape(pd_tensor, phi::IntArray(reshape_vec));
paddle::Tensor result = paddle::experimental::expand(
reshaped, phi::IntArray(target_size_vec));
return Tensor(result);
}
PD_THROW("expand(): the expanded size of the tensor (",
target_size_vec[fail_dim],
") must match the existing size (",
reshape_vec[fail_dim],
") at non-singleton dimension ",
fail_dim,
".");
} else if (input_rank == target_rank) {
// Same rank - check if we can use expand directly
bool can_use_paddle_expand = true;
size_t fail_dim = 0;
for (size_t i = 0; i < target_rank; ++i) {
auto in_size = input_dims[i];
auto target_size = target_size_vec[i];
if (in_size != 1 && in_size != target_size) {
can_use_paddle_expand = false;
fail_dim = i;
break;
}
}
if (can_use_paddle_expand) {
paddle::Tensor result = paddle::experimental::expand(
pd_tensor, phi::IntArray(target_size_vec));
return Tensor(result);
}
PD_THROW("expand(): the expanded size of the tensor (",
target_size_vec[fail_dim],
") must match the existing size (",
input_dims[fail_dim],
") at non-singleton dimension ",
fail_dim,
".");
} else {
PD_THROW("expand(): the number of sizes provided (",
target_rank,
") must be greater or equal to the number of dimensions in the "
"tensor (",
input_rank,
").");
}
}
} // namespace at
namespace at {
// Member function: Tensor::expand
inline Tensor Tensor::expand(at::IntArrayRef size, bool implicit) const {
return at::expand(*this, size, implicit);
}
} // namespace at
@@ -0,0 +1,26 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <ATen/ops/expand.h>
namespace at {
inline Tensor Tensor::expand_as(const Tensor& other) const {
return at::expand(*this, other.sizes());
}
} // namespace at
@@ -0,0 +1,106 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/core/TensorOptions.h>
#include <utils/pinned_place.h>
#include <optional>
#include "paddle/phi/api/include/api.h"
#include "paddle/phi/common/place.h"
namespace at {
// eye(n) — n×n identity matrix
inline at::Tensor eye(int64_t n, at::TensorOptions options = {}) {
if (options.pinned_memory()) {
// Pinning memory is only supported for CPU tensors
if (options.has_device() && !options.device().is_cpu()) {
PD_THROW(
"pin_memory=true requires device to be CPU, but got non-CPU device");
}
phi::Place base_place = options._PD_GetPlace();
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(base_place);
auto dense = paddle::experimental::eye(
n,
/*num_columns=*/-1,
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
phi::CPUPlace());
return dense.copy_to(pinned_place, /*blocking=*/true);
}
return paddle::experimental::eye(
n,
/*num_columns=*/-1,
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
options._PD_GetPlace());
}
// eye(n, m) — n×m identity-like matrix
inline at::Tensor eye(int64_t n, int64_t m, at::TensorOptions options = {}) {
if (options.pinned_memory()) {
// Pinning memory is only supported for CPU tensors
if (options.has_device() && !options.device().is_cpu()) {
PD_THROW(
"pin_memory=true requires device to be CPU, but got non-CPU device");
}
phi::Place base_place = options._PD_GetPlace();
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(base_place);
auto dense = paddle::experimental::eye(
n,
m,
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
phi::CPUPlace());
return dense.copy_to(pinned_place, /*blocking=*/true);
}
return paddle::experimental::eye(
n,
m,
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
options._PD_GetPlace());
}
// eye(n, dtype, layout, device, pin_memory)
inline at::Tensor eye(int64_t n,
::std::optional<at::ScalarType> dtype,
::std::optional<at::Layout> layout,
::std::optional<at::Device> device,
::std::optional<bool> pin_memory) {
PD_CHECK(!layout.has_value(), "`layout` is not supported now.");
auto options =
at::TensorOptions()
.dtype(dtype.value_or(c10::get_default_dtype_as_scalartype()))
.device(device.value_or(at::kCPU))
.pinned_memory(pin_memory);
return eye(n, options);
}
// eye(n, m, dtype, layout, device, pin_memory)
inline at::Tensor eye(int64_t n,
int64_t m,
::std::optional<at::ScalarType> dtype,
::std::optional<at::Layout> layout,
::std::optional<at::Device> device,
::std::optional<bool> pin_memory) {
PD_CHECK(!layout.has_value(), "`layout` is not supported now.");
auto options =
at::TensorOptions()
.dtype(dtype.value_or(c10::get_default_dtype_as_scalartype()))
.device(device.value_or(at::kCPU))
.pinned_memory(pin_memory);
return eye(n, m, options);
}
} // namespace at
@@ -0,0 +1,37 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
namespace at {
inline at::Tensor flatten(const at::Tensor& self,
int64_t start_dim = 0,
int64_t end_dim = -1) {
return Tensor(paddle::experimental::flatten(self._PD_GetInner(),
static_cast<int>(start_dim),
static_cast<int>(end_dim)));
}
} // namespace at
namespace at {
inline at::Tensor Tensor::flatten(int64_t start_dim, int64_t end_dim) const {
return at::flatten(*this, start_dim, end_dim);
}
} // namespace at
@@ -0,0 +1,219 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include "paddle/phi/api/include/tensor_utils.h"
namespace at {
namespace detail {
inline void noopDelete(void* /*unused*/) {}
} // namespace detail
class TensorMaker {
friend TensorMaker for_blob(void* data, IntArrayRef sizes) noexcept;
public:
using ContextDeleter = DeleterFnPtr;
TensorMaker& strides(OptionalIntArrayRef value) noexcept {
strides_ = value;
return *this;
}
TensorMaker& storage_offset(std::optional<int64_t> value) noexcept {
storage_offset_ = value;
return *this;
}
TensorMaker& deleter(std::function<void(void*)> value) noexcept {
deleter_ = std::move(value);
return *this;
}
TensorMaker& context(void* value, ContextDeleter deleter = nullptr) noexcept {
ctx_ = std::unique_ptr<void, ContextDeleter>{
value, deleter != nullptr ? deleter : detail::noopDelete};
return *this;
}
TensorMaker& target_device(std::optional<Device> value) noexcept {
device_ = value;
return *this;
}
TensorMaker& options(TensorOptions value) noexcept {
opts_ = value;
return *this;
}
TensorMaker& resizeable_storage() noexcept {
resizeable_ = true;
return *this;
}
Tensor make_tensor() {
PD_CHECK(!deleter_ || !ctx_,
"The deleter and context arguments are mutually exclusive.");
PD_CHECK(!storage_offset_.has_value() || storage_offset_.value() == 0,
"storage_offset` should be zero.");
if (device_.has_value() && opts_.has_device() &&
opts_.device().has_index()) {
PD_CHECK(opts_.device() == *device_,
"Specified device ",
opts_.device(),
" does not match device of data ",
*device_);
}
phi::Place pd_place;
if (device_.has_value()) {
pd_place = device_->_PD_GetInner();
} else if (opts_.has_device() && opts_.device().has_index()) {
pd_place = opts_.device()._PD_GetInner();
} else {
pd_place = phi::Place(); // UNDEFINED → auto-detect inside from_blob
}
// Build paddle deleter: prefer explicit deleter_, then wrap ctx_ so its
// lifetime is tied to the tensor allocation.
paddle::Deleter pd_deleter = nullptr;
if (deleter_) {
pd_deleter = deleter_;
} else if (ctx_) {
// shared_ptr takes ownership of the context and calls its deleter when
// the last copy (held in the lambda) is destroyed.
auto shared_ctx =
std::shared_ptr<void>(ctx_.release(), ctx_.get_deleter());
pd_deleter = [shared_ctx](void* /*data*/) {};
}
if (strides_.has_value()) {
return paddle::from_blob(
data_,
sizes_._PD_ToPaddleIntArray(),
strides_.value()._PD_ToPaddleIntArray(),
compat::_PD_AtenScalarTypeToPhiDataType(opts_.dtype()),
phi::DataLayout::NCHW,
pd_place,
pd_deleter);
} else {
return paddle::from_blob(
data_,
sizes_._PD_ToPaddleIntArray(),
compat::_PD_AtenScalarTypeToPhiDataType(opts_.dtype()),
phi::DataLayout::NCHW,
pd_place,
pd_deleter);
}
}
private:
explicit TensorMaker(void* data, IntArrayRef sizes) noexcept
: data_{data}, sizes_{sizes} {}
std::size_t computeStorageSize() const noexcept;
DataPtr makeDataPtrFromDeleter() noexcept;
DataPtr makeDataPtrFromContext() noexcept;
IntArrayRef makeTempSizes() const noexcept;
void* data_;
IntArrayRef sizes_;
OptionalIntArrayRef strides_;
std::optional<int64_t> storage_offset_;
std::function<void(void*)> deleter_;
std::unique_ptr<void, ContextDeleter> ctx_{nullptr, detail::noopDelete};
std::optional<Device> device_;
TensorOptions opts_;
bool resizeable_{};
};
inline TensorMaker for_blob(void* data, IntArrayRef sizes) noexcept {
return TensorMaker{data, sizes};
}
inline Tensor from_blob(
void* data,
IntArrayRef sizes,
IntArrayRef strides,
const std::function<void(void*)>& deleter,
const TensorOptions& options = {},
const std::optional<Device> target_device = std::nullopt) {
return for_blob(data, sizes)
.strides(strides)
.deleter(deleter)
.options(options)
.target_device(target_device)
.make_tensor();
}
inline Tensor from_blob(
void* data,
IntArrayRef sizes,
IntArrayRef strides,
int64_t storage_offset,
const std::function<void(void*)>& deleter,
const TensorOptions& options = {},
const std::optional<Device> target_device = std::nullopt) {
return for_blob(data, sizes)
.strides(strides)
.storage_offset(storage_offset)
.deleter(deleter)
.options(options)
.target_device(target_device)
.make_tensor();
}
inline Tensor from_blob(
void* data,
IntArrayRef sizes,
std::function<void(void*)> deleter,
const TensorOptions& options = {},
const std::optional<Device> target_device = std::nullopt) {
return for_blob(data, sizes)
.deleter(std::move(deleter))
.options(options)
.target_device(target_device)
.make_tensor();
}
inline Tensor from_blob(void* data,
IntArrayRef sizes,
IntArrayRef strides,
const TensorOptions& options = {}) {
return for_blob(data, sizes).strides(strides).options(options).make_tensor();
}
inline Tensor from_blob(void* data,
IntArrayRef sizes,
const TensorOptions& options = {}) {
return for_blob(data, sizes).options(options).make_tensor();
}
} // namespace at
@@ -0,0 +1,104 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/core/SymIntArrayRef.h>
#include <c10/core/TensorOptions.h>
#include <utils/pinned_place.h>
#include <optional>
#include <string_view>
#include "paddle/phi/api/include/api.h"
#include "paddle/phi/common/place.h"
namespace at {
inline at::Tensor full(at::IntArrayRef size,
const at::Scalar& fill_value,
at::TensorOptions options = {}) {
if (options.pinned_memory()) {
// Pinning memory is only supported for CPU tensors
if (options.has_device() && !options.device().is_cpu()) {
PD_THROW(
"pin_memory=true requires device to be CPU, but got non-CPU device");
}
phi::Place base_place = options._PD_GetPlace();
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(base_place);
auto dense = paddle::experimental::full(
size._PD_ToPaddleIntArray(),
fill_value,
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
phi::CPUPlace());
return dense.copy_to(pinned_place, /*blocking=*/true);
}
return paddle::experimental::full(
size._PD_ToPaddleIntArray(),
fill_value,
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
options._PD_GetPlace());
}
inline at::Tensor full(at::IntArrayRef size,
const at::Scalar& fill_value,
::std::optional<at::ScalarType> dtype,
::std::optional<at::Layout> layout,
::std::optional<at::Device> device,
::std::optional<bool> pin_memory) {
PD_CHECK(!layout.has_value(), "`layout` is not supported now.");
auto options =
at::TensorOptions()
.dtype(dtype.value_or(c10::get_default_dtype_as_scalartype()))
.device(device.value_or(at::kCPU))
.pinned_memory(pin_memory);
return full(size, fill_value, options);
}
inline at::Tensor full_symint(c10::SymIntArrayRef size,
const at::Scalar& fill_value,
at::TensorOptions options = {}) {
if (options.pinned_memory()) {
phi::Place base_place = options._PD_GetPlace();
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(base_place);
auto dense = paddle::experimental::full(
size._PD_ToPaddleIntArray(),
fill_value,
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
phi::CPUPlace());
return dense.copy_to(pinned_place, /*blocking=*/true);
}
return paddle::experimental::full(
size._PD_ToPaddleIntArray(),
fill_value,
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
options._PD_GetPlace());
}
inline at::Tensor full_symint(c10::SymIntArrayRef size,
const at::Scalar& fill_value,
::std::optional<at::ScalarType> dtype,
::std::optional<at::Layout> layout,
::std::optional<at::Device> device,
::std::optional<bool> pin_memory) {
PD_CHECK(!layout.has_value(), "`layout` is not supported now.");
auto options =
at::TensorOptions()
.dtype(dtype.value_or(c10::get_default_dtype_as_scalartype()))
.device(device.value_or(at::kCPU))
.pinned_memory(pin_memory);
return full_symint(size, fill_value, options);
}
} // namespace at
@@ -0,0 +1,47 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <ATen/ops/tensor_split.h>
namespace at {
inline std::vector<at::Tensor> hsplit(const at::Tensor& self,
int64_t sections) {
// For 1D tensors, split along dim 0; otherwise split along dim 1
int64_t dim = self._PD_GetInner().dims().size() == 1 ? 0 : 1;
return at::tensor_split(self, sections, dim);
}
inline std::vector<at::Tensor> hsplit(const at::Tensor& self,
at::IntArrayRef indices) {
int64_t dim = self._PD_GetInner().dims().size() == 1 ? 0 : 1;
return at::tensor_split(self, indices, dim);
}
} // namespace at
namespace at {
inline std::vector<at::Tensor> Tensor::hsplit(int64_t sections) const {
return at::hsplit(*this, sections);
}
inline std::vector<at::Tensor> Tensor::hsplit(at::IntArrayRef indices) const {
return at::hsplit(*this, indices);
}
} // namespace at
@@ -0,0 +1,203 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/TensorIndexing.h>
#include <ATen/core/Tensor.h>
#include <c10/core/List.h>
namespace at::indexing {
inline TensorIndex::TensorIndex(const at::Tensor& tensor)
: tensor_(std::make_shared<at::Tensor>(tensor)),
type_(TensorIndexType::Tensor) {}
inline const at::Tensor& TensorIndex::tensor() const { return *tensor_; }
} // namespace at::indexing
namespace at::detail {
inline bool _PD_is_full_slice(const at::indexing::Slice& slice) {
return static_cast<int64_t>(slice.start()) == 0 &&
static_cast<int64_t>(slice.stop()) == at::indexing::INDEX_MAX &&
static_cast<int64_t>(slice.step()) == 1;
}
inline at::Tensor _PD_apply_tensor_index(
const at::Tensor& self, ArrayRef<at::indexing::TensorIndex> indices) {
int64_t output_dim = 0;
int tensor_index_count = 0;
at::Tensor result = self;
for (const auto& index : indices) {
if (index.is_tensor()) {
++tensor_index_count;
PD_CHECK(tensor_index_count == 1,
"Multiple tensor indices mixed with None/Slice are not "
"supported yet.");
result = paddle::experimental::index_select(
result._PD_GetInner(), index.tensor()._PD_GetInner(), output_dim);
++output_dim;
} else if (index.is_none()) {
result =
paddle::experimental::unsqueeze(result._PD_GetInner(), {output_dim});
++output_dim;
} else if (index.is_slice()) {
const auto& slice = index.slice();
PD_CHECK(_PD_is_full_slice(slice),
"Only full Slice() is supported when mixed with tensor/None "
"indices.");
++output_dim;
}
}
return result;
}
inline at::Tensor _PD_index_tensor_indices(
const at::Tensor& self, ArrayRef<at::indexing::TensorIndex> indices) {
if (indices.size() == 0) {
PD_THROW("index() cannot be called with an empty index list");
}
bool has_slice = false;
bool has_tensor_like = false;
for (const auto& index : indices) {
has_slice = has_slice || index.is_slice();
has_tensor_like = has_tensor_like || index.is_tensor() || index.is_none();
PD_CHECK(!index.is_ellipsis(), "Ellipsis index is not supported yet.");
PD_CHECK(!index.is_integer(), "Integer index is not supported yet.");
PD_CHECK(!index.is_boolean(), "Boolean index is not supported yet.");
}
if (has_slice && !has_tensor_like) {
std::vector<int64_t> axes;
std::vector<int64_t> starts;
std::vector<int64_t> ends;
std::vector<int64_t> strides;
axes.reserve(indices.size());
starts.reserve(indices.size());
ends.reserve(indices.size());
strides.reserve(indices.size());
int64_t dim = 0;
for (const auto& index : indices) {
const auto& slice = index.slice();
axes.push_back(dim++);
starts.push_back(static_cast<int64_t>(slice.start()));
ends.push_back(static_cast<int64_t>(slice.stop()));
strides.push_back(static_cast<int64_t>(slice.step()));
}
return paddle::experimental::slice(
self._PD_GetInner(), axes, starts, ends, strides, {});
}
if (has_slice) {
return _PD_apply_tensor_index(self, indices);
}
c10::List<::std::optional<at::Tensor>> tensor_indices;
for (const auto& index : indices) {
if (index.is_none()) {
tensor_indices.push_back(::std::nullopt);
} else if (index.is_tensor()) {
tensor_indices.push_back(index.tensor());
}
}
return self.index(tensor_indices);
}
} // namespace at::detail
namespace at {
inline at::Tensor index(const at::Tensor& self,
const c10::List<::std::optional<at::Tensor>>& indices) {
if (indices.empty()) {
return self;
}
bool all_none = true;
for (const auto& idx : indices) {
if (idx.has_value()) {
all_none = false;
break;
}
}
if (all_none) {
return self;
}
std::vector<paddle::Tensor> pd_indices;
std::vector<bool> has_index(indices.size(), false);
pd_indices.reserve(indices.size());
for (size_t i = 0; i < indices.size(); ++i) {
if (indices[i].has_value()) {
pd_indices.push_back(indices[i].value()._PD_GetInner());
has_index[i] = true;
} else {
pd_indices.push_back(paddle::Tensor());
}
}
int non_none_count = 0;
size_t first_non_none = 0;
for (size_t i = 0; i < has_index.size(); ++i) {
if (has_index[i]) {
non_none_count++;
first_non_none = i;
}
}
if (non_none_count == 1 && first_non_none == 0) {
return paddle::experimental::index_select(
self._PD_GetInner(), pd_indices[0], 0);
}
if (non_none_count == static_cast<int>(indices.size())) {
auto stacked_indices = paddle::experimental::stack(pd_indices, -1);
return paddle::experimental::gather_nd(self._PD_GetInner(),
stacked_indices);
}
auto self_dims = self._PD_GetInner().dims();
int self_rank = self_dims.size();
at::Tensor result = self;
for (size_t i = 0; i < indices.size() && i < static_cast<size_t>(self_rank);
++i) {
if (has_index[i]) {
paddle::Tensor pd_result = result._PD_GetInner();
result = paddle::experimental::index_select(
pd_result, pd_indices[i], static_cast<int>(i));
}
}
return result;
}
} // namespace at
namespace at {
inline at::Tensor Tensor::index(
ArrayRef<at::indexing::TensorIndex> indices) const {
return at::detail::_PD_index_tensor_indices(*this, indices);
}
} // namespace at
@@ -0,0 +1,205 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <ATen/ops/index.h>
#include <c10/core/List.h>
#include <c10/core/Scalar.h>
#include <vector>
#include "paddle/phi/api/include/api.h"
namespace at::detail {
inline std::vector<at::Tensor> _PD_convert_indices_list(
const c10::List<::std::optional<at::Tensor>>& indices) {
std::vector<at::Tensor> result;
result.reserve(indices.size());
for (const auto& idx : indices) {
if (idx.has_value()) {
result.push_back(idx.value());
}
}
return result;
}
inline c10::List<::std::optional<at::Tensor>> _PD_convert_tensor_index_list(
ArrayRef<at::indexing::TensorIndex> indices) {
c10::List<::std::optional<at::Tensor>> result;
for (const auto& index : indices) {
PD_CHECK(!index.is_ellipsis(), "Ellipsis index is not supported yet.");
PD_CHECK(!index.is_integer(), "Integer index is not supported yet.");
PD_CHECK(!index.is_boolean(), "Boolean index is not supported yet.");
if (index.is_slice()) {
PD_CHECK(_PD_is_full_slice(index.slice()),
"Only full Slice() is supported in index_put_ TensorIndex "
"paths.");
} else if (index.is_tensor()) {
result.push_back(index.tensor());
}
}
return result;
}
inline at::Tensor _PD_squeeze_newaxis_value(
const at::Tensor& values, ArrayRef<at::indexing::TensorIndex> indices) {
std::vector<int64_t> value_shape(values.sizes().begin(),
values.sizes().end());
size_t value_dim = 0;
bool changed = false;
for (const auto& index : indices) {
if (index.is_none()) {
if (!value_shape.empty()) {
PD_CHECK(value_dim < value_shape.size(),
"index_put_ value rank is too small for None index.");
PD_CHECK(value_shape[value_dim] == 1,
"index_put_ expected value dimension inserted by None to "
"have size 1, but got ",
value_shape[value_dim],
".");
value_shape.erase(value_shape.begin() + value_dim);
changed = true;
}
} else if (index.is_tensor()) {
if (!value_shape.empty()) {
++value_dim;
}
} else if (index.is_slice()) {
PD_CHECK(_PD_is_full_slice(index.slice()),
"Only full Slice() is supported in index_put_ TensorIndex "
"paths.");
if (!value_shape.empty()) {
++value_dim;
}
} else {
PD_CHECK(!index.is_ellipsis(), "Ellipsis index is not supported yet.");
PD_CHECK(!index.is_integer(), "Integer index is not supported yet.");
PD_CHECK(!index.is_boolean(), "Boolean index is not supported yet.");
}
}
if (!changed) {
return values;
}
return paddle::experimental::reshape(values._PD_GetInner(),
phi::IntArray(value_shape));
}
} // namespace at::detail
namespace at {
// index_put_: Set values at specified indices (in-place)
inline at::Tensor& index_put_(
at::Tensor& self, // NOLINT(runtime/references)
const c10::List<::std::optional<at::Tensor>>& indices,
const at::Tensor& values,
bool accumulate = false) {
std::vector<paddle::Tensor> pd_indices;
pd_indices.reserve(indices.size());
for (const auto& idx : indices) {
if (idx.has_value()) {
pd_indices.push_back(idx.value()._PD_GetInner());
}
}
paddle::experimental::index_put_(
self._PD_GetInner(), pd_indices, values._PD_GetInner(), accumulate);
return self;
}
// index_put: Non-inplace version
inline at::Tensor index_put(
const at::Tensor& self,
const c10::List<::std::optional<at::Tensor>>& indices,
const at::Tensor& values,
bool accumulate = false) {
std::vector<paddle::Tensor> pd_indices;
pd_indices.reserve(indices.size());
for (const auto& idx : indices) {
if (idx.has_value()) {
pd_indices.push_back(idx.value()._PD_GetInner());
}
}
return paddle::experimental::index_put(
self._PD_GetInner(), pd_indices, values._PD_GetInner(), accumulate);
}
} // namespace at
namespace at {
inline at::Tensor Tensor::index(
const c10::List<::std::optional<at::Tensor>>& indices) const {
return at::index(*this, indices);
}
inline at::Tensor& Tensor::index_put_(
const c10::List<::std::optional<at::Tensor>>& indices,
const at::Tensor& values,
bool accumulate) const {
return at::index_put_(
const_cast<at::Tensor&>(*this), indices, values, accumulate);
}
inline at::Tensor& Tensor::index_put_(
ArrayRef<at::indexing::TensorIndex> indices, Tensor const& rhs) {
auto tensor_indices = detail::_PD_convert_tensor_index_list(indices);
at::Tensor values = detail::_PD_squeeze_newaxis_value(rhs, indices);
if (tensor_indices.empty()) {
return copy_(values);
}
return index_put_(tensor_indices, values);
}
inline at::Tensor& Tensor::index_put_(
ArrayRef<at::indexing::TensorIndex> indices, const Scalar& v) {
auto tensor_indices = detail::_PD_convert_tensor_index_list(indices);
if (tensor_indices.empty()) {
std::vector<int64_t> value_shape(this->sizes().begin(),
this->sizes().end());
auto scalar_tensor =
at::Tensor(paddle::experimental::full(phi::IntArray(value_shape),
phi::Scalar(v.to<double>()),
this->_PD_GetInner().dtype()));
return copy_(scalar_tensor);
}
auto scalar_tensor = at::Tensor(paddle::experimental::full(
{}, phi::Scalar(v.to<double>()), this->_PD_GetInner().dtype()));
return index_put_(indices, scalar_tensor);
}
inline at::Tensor& Tensor::index_put_(
std::initializer_list<at::indexing::TensorIndex> indices,
Tensor const& rhs) {
return index_put_(ArrayRef<at::indexing::TensorIndex>(indices), rhs);
}
inline at::Tensor& Tensor::index_put_(
std::initializer_list<at::indexing::TensorIndex> indices, const Scalar& v) {
return index_put_(ArrayRef<at::indexing::TensorIndex>(indices), v);
}
inline at::Tensor Tensor::index_put(
const c10::List<::std::optional<at::Tensor>>& indices,
const at::Tensor& values,
bool accumulate) const {
return at::index_put(*this, indices, values, accumulate);
}
} // namespace at
@@ -0,0 +1,33 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include "paddle/phi/core/sparse_coo_tensor.h"
namespace at {} // namespace at
namespace at {
inline bool Tensor::is_coalesced() const {
PD_CHECK(tensor_.layout() == common::DataLayout::SPARSE_COO,
"is_coalesced expected sparse coordinate tensor layout but got ",
layout());
auto sparse_coo_tensor =
std::dynamic_pointer_cast<phi::SparseCooTensor>(tensor_.impl());
return sparse_coo_tensor->coalesced();
}
} // namespace at
@@ -0,0 +1,44 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <ATen/ops/_local_scalar_dense.h>
namespace at {} // namespace at
namespace at {
inline at::Scalar Tensor::item() const {
auto numel = this->sym_numel();
PD_CHECK(numel == 1,
"a Tensor with ",
numel,
" elements cannot be converted to Scalar");
if (this->is_sparse()) {
if (this->_nnz() == 0) return Scalar(0);
if (this->is_coalesced()) return at::_local_scalar_dense(this->_values());
return at::_local_scalar_dense(this->_values().sum());
} else {
return _local_scalar_dense(*this);
}
}
template <typename T>
T Tensor::item() const {
return item().to<T>();
}
} // namespace at
@@ -0,0 +1,35 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
namespace at {
inline at::Tensor masked_select(const at::Tensor& self,
const at::Tensor& mask) {
return Tensor(paddle::experimental::masked_select(self._PD_GetInner(),
mask._PD_GetInner()));
}
} // namespace at
namespace at {
inline at::Tensor Tensor::masked_select(const at::Tensor& mask) const {
return at::masked_select(*this, mask);
}
} // namespace at
@@ -0,0 +1,115 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
namespace at {
inline at::Tensor narrow(const at::Tensor& self,
int64_t dim,
int64_t start,
int64_t length) {
// Bounds checks matching PyTorch behavior
PD_CHECK(self.dim() > 0, "narrow() cannot be applied to a 0-dim tensor.");
PD_CHECK(length >= 0, "narrow(): length must be non-negative.");
// Normalize negative dim
int64_t ndim = self.dim();
if (dim < 0) dim += ndim;
PD_CHECK(dim >= 0 && dim < ndim,
"start out of range (expected to be in range of [",
-ndim,
", ",
ndim - 1,
"], but got ",
dim,
")");
int64_t cur_size = self.sizes()[dim];
// Wrap negative start (matching PyTorch: only wrap when start != cur_size)
if (start < 0) {
start = start + cur_size;
}
PD_CHECK(start <= cur_size - length,
"start (",
start,
") + length (",
length,
") exceeds dimension size (",
cur_size,
").");
// Use slice to implement narrow: narrow(dim, start, length) is equivalent
// to slice(dim, start, start + length)
return Tensor(paddle::experimental::slice(
self._PD_GetInner(), {dim}, {start}, {start + length}, {1}, {}));
}
inline at::Tensor narrow_symint(const at::Tensor& self,
int64_t dim,
c10::SymInt start,
c10::SymInt length) {
return narrow(self, dim, start, length);
}
inline at::Tensor narrow(const at::Tensor& self,
int64_t dim,
const at::Tensor& start,
int64_t length) {
// Extract scalar value from start tensor
PD_CHECK(start.numel() == 1,
"start must be a 0-dim tensor or 1-element tensor");
int64_t start_val = start.item<int64_t>();
return narrow(self, dim, start_val, length);
}
inline at::Tensor narrow_symint(const at::Tensor& self,
int64_t dim,
const at::Tensor& start,
c10::SymInt length) {
return narrow(self, dim, start, length);
}
} // namespace at
namespace at {
inline at::Tensor Tensor::narrow(int64_t dim,
int64_t start,
int64_t length) const {
return at::narrow(*this, dim, start, length);
}
inline at::Tensor Tensor::narrow_symint(int64_t dim,
c10::SymInt start,
c10::SymInt length) const {
return at::narrow_symint(*this, dim, start, length);
}
inline at::Tensor Tensor::narrow(int64_t dim,
const at::Tensor& start,
int64_t length) const {
return at::narrow(*this, dim, start, length);
}
inline at::Tensor Tensor::narrow_symint(int64_t dim,
const at::Tensor& start,
c10::SymInt length) const {
return at::narrow_symint(*this, dim, start, length);
}
} // namespace at
@@ -0,0 +1,53 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <ATen/ops/narrow.h>
namespace at {
inline at::Tensor narrow_copy(const at::Tensor& self,
int64_t dim,
int64_t start,
int64_t length) {
// narrow_copy returns a copy of the narrowed tensor
return narrow(self, dim, start, length).clone();
}
inline at::Tensor narrow_copy_symint(const at::Tensor& self,
int64_t dim,
c10::SymInt start,
c10::SymInt length) {
return narrow_copy(self, dim, start, length);
}
} // namespace at
namespace at {
inline at::Tensor Tensor::narrow_copy(int64_t dim,
int64_t start,
int64_t length) const {
return at::narrow_copy(*this, dim, start, length);
}
inline at::Tensor Tensor::narrow_copy_symint(int64_t dim,
c10::SymInt start,
c10::SymInt length) const {
return at::narrow_copy_symint(*this, dim, start, length);
}
} // namespace at
@@ -0,0 +1,68 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/core/TensorOptions.h>
#include <utils/pinned_place.h>
#include <optional>
#include <string_view>
#include "paddle/phi/api/include/api.h"
#include "paddle/phi/common/place.h"
namespace at {
// Member function: Tensor::new_empty
inline Tensor Tensor::new_empty(at::IntArrayRef size,
at::TensorOptions options) const {
caffe2::TypeMeta actual_dtype = options.dtype_opt().value_or(dtype());
auto actual_device = options.device_opt().value_or(device());
auto actual_pin_memory = options.pinned_memory();
auto pd_dtype = compat::_PD_AtenScalarTypeToPhiDataType(actual_dtype);
auto pd_place = actual_device._PD_GetInner();
paddle::Tensor result;
if (actual_pin_memory) {
// Pinning memory is only supported for CPU tensors
if (options.has_device() && !actual_device.is_cpu()) {
PD_THROW(
"pin_memory=true requires device to be CPU, but got non-CPU device");
}
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(pd_place);
auto dense_cpu = paddle::experimental::empty(
size._PD_ToPaddleIntArray(), pd_dtype, phi::CPUPlace());
result = dense_cpu.copy_to(pinned_place, /*blocking=*/true);
} else {
result = paddle::experimental::empty(
size._PD_ToPaddleIntArray(), pd_dtype, pd_place);
}
return Tensor(result);
}
inline Tensor Tensor::new_empty(at::IntArrayRef size,
::std::optional<at::ScalarType> dtype,
::std::optional<at::Layout>,
::std::optional<at::Device> device,
::std::optional<bool> pin_memory) const {
auto options = at::TensorOptions()
.dtype(dtype.value_or(this->scalar_type()))
.device(device.value_or(this->device()))
.pinned_memory(pin_memory);
return new_empty(size, options);
}
} // namespace at
@@ -0,0 +1,70 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/core/TensorOptions.h>
#include <utils/pinned_place.h>
#include <optional>
#include <string_view>
#include "paddle/phi/api/include/api.h"
#include "paddle/phi/common/place.h"
namespace at {
// Member function: Tensor::new_full
inline Tensor Tensor::new_full(at::IntArrayRef size,
const at::Scalar& fill_value,
at::TensorOptions options) const {
caffe2::TypeMeta actual_dtype = options.dtype_opt().value_or(dtype());
auto actual_device = options.device_opt().value_or(device());
auto actual_pin_memory = options.pinned_memory();
auto pd_dtype = compat::_PD_AtenScalarTypeToPhiDataType(actual_dtype);
auto pd_place = actual_device._PD_GetInner();
paddle::Tensor result;
if (actual_pin_memory) {
// Pinning memory is only supported for CPU tensors
if (options.has_device() && !actual_device.is_cpu()) {
PD_THROW(
"pin_memory=true requires device to be CPU, but got non-CPU device");
}
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(pd_place);
auto dense_cpu = paddle::experimental::full(
size._PD_ToPaddleIntArray(), fill_value, pd_dtype, phi::CPUPlace());
result = dense_cpu.copy_to(pinned_place, /*blocking=*/true);
} else {
result = paddle::experimental::full(
size._PD_ToPaddleIntArray(), fill_value, pd_dtype, pd_place);
}
return Tensor(result);
}
inline Tensor Tensor::new_full(at::IntArrayRef size,
const at::Scalar& fill_value,
::std::optional<at::ScalarType> dtype,
::std::optional<at::Layout>,
::std::optional<at::Device> device,
::std::optional<bool> pin_memory) const {
auto options = at::TensorOptions()
.dtype(dtype.value_or(this->scalar_type()))
.device(device.value_or(this->device()))
.pinned_memory(pin_memory);
return new_full(size, fill_value, options);
}
} // namespace at
@@ -0,0 +1,68 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/core/TensorOptions.h>
#include <utils/pinned_place.h>
#include <optional>
#include <string_view>
#include "paddle/phi/api/include/api.h"
#include "paddle/phi/common/place.h"
namespace at {
// Member function: Tensor::new_ones
inline Tensor Tensor::new_ones(at::IntArrayRef size,
at::TensorOptions options) const {
caffe2::TypeMeta actual_dtype = options.dtype_opt().value_or(dtype());
auto actual_device = options.device_opt().value_or(device());
auto actual_pin_memory = options.pinned_memory();
auto pd_dtype = compat::_PD_AtenScalarTypeToPhiDataType(actual_dtype);
auto pd_place = actual_device._PD_GetInner();
paddle::Tensor result;
if (actual_pin_memory) {
// Pinning memory is only supported for CPU tensors
if (options.has_device() && !actual_device.is_cpu()) {
PD_THROW(
"pin_memory=true requires device to be CPU, but got non-CPU device");
}
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(pd_place);
auto dense_cpu = paddle::experimental::ones(
size._PD_ToPaddleIntArray(), pd_dtype, phi::CPUPlace());
result = dense_cpu.copy_to(pinned_place, /*blocking=*/true);
} else {
result = paddle::experimental::ones(
size._PD_ToPaddleIntArray(), pd_dtype, pd_place);
}
return Tensor(result);
}
inline Tensor Tensor::new_ones(at::IntArrayRef size,
::std::optional<at::ScalarType> dtype,
::std::optional<at::Layout>,
::std::optional<at::Device> device,
::std::optional<bool> pin_memory) const {
auto options = at::TensorOptions()
.dtype(dtype.value_or(this->scalar_type()))
.device(device.value_or(this->device()))
.pinned_memory(pin_memory);
return new_ones(size, options);
}
} // namespace at
@@ -0,0 +1,69 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/core/SymIntArrayRef.h>
#include <c10/core/TensorOptions.h>
#include <utils/pinned_place.h>
#include <optional>
#include <string_view>
#include "paddle/phi/api/include/api.h"
#include "paddle/phi/common/place.h"
namespace at {
// Member function: Tensor::new_zeros
inline Tensor Tensor::new_zeros(at::IntArrayRef size,
at::TensorOptions options) const {
caffe2::TypeMeta actual_dtype = options.dtype_opt().value_or(dtype());
auto actual_device = options.device_opt().value_or(device());
auto actual_pin_memory = options.pinned_memory();
auto pd_dtype = compat::_PD_AtenScalarTypeToPhiDataType(actual_dtype);
auto pd_place = actual_device._PD_GetInner();
paddle::Tensor result;
if (actual_pin_memory) {
// Pinning memory is only supported for CPU tensors
if (options.has_device() && !actual_device.is_cpu()) {
PD_THROW(
"pin_memory=true requires device to be CPU, but got non-CPU device");
}
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(pd_place);
auto dense_cpu = paddle::experimental::zeros(
size._PD_ToPaddleIntArray(), pd_dtype, phi::CPUPlace());
result = dense_cpu.copy_to(pinned_place, /*blocking=*/true);
} else {
result = paddle::experimental::zeros(
size._PD_ToPaddleIntArray(), pd_dtype, pd_place);
}
return Tensor(result);
}
inline Tensor Tensor::new_zeros(at::IntArrayRef size,
::std::optional<at::ScalarType> dtype,
::std::optional<at::Layout>,
::std::optional<at::Device> device,
::std::optional<bool> pin_memory) const {
auto options = at::TensorOptions()
.dtype(dtype.value_or(this->scalar_type()))
.device(device.value_or(this->device()))
.pinned_memory(pin_memory);
return new_zeros(size, options);
}
} // namespace at
@@ -0,0 +1,99 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/core/TensorOptions.h>
#include <utils/pinned_place.h>
#include <optional>
#include <string_view>
#include "paddle/phi/api/include/api.h"
#include "paddle/phi/common/place.h"
namespace at {
inline at::Tensor ones(at::IntArrayRef size, at::TensorOptions options = {}) {
if (options.pinned_memory()) {
// Pinning memory is only supported for CPU tensors
if (options.has_device() && !options.device().is_cpu()) {
PD_THROW(
"pin_memory=true requires device to be CPU, but got non-CPU device");
}
phi::Place base_place = options._PD_GetPlace();
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(base_place);
auto dense = paddle::experimental::ones(
size._PD_ToPaddleIntArray(),
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
phi::CPUPlace());
return dense.copy_to(pinned_place, /*blocking=*/true);
}
return paddle::experimental::ones(
size._PD_ToPaddleIntArray(),
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
options._PD_GetPlace());
}
inline at::Tensor ones(at::IntArrayRef size,
::std::optional<at::ScalarType> dtype,
::std::optional<at::Layout> layout,
::std::optional<at::Device> device,
::std::optional<bool> pin_memory) {
PD_CHECK(!layout.has_value(), "`layout` is not supported now.");
auto options =
at::TensorOptions()
.dtype(dtype.value_or(c10::get_default_dtype_as_scalartype()))
.device(device.value_or(at::kCPU))
.pinned_memory(pin_memory);
return ones(size, options);
}
inline at::Tensor ones_symint(c10::SymIntArrayRef size,
at::TensorOptions options = {}) {
if (options.pinned_memory()) {
// Pinning memory is only supported for CPU tensors
if (options.has_device() && !options.device().is_cpu()) {
PD_THROW(
"pin_memory=true requires device to be CPU, but got non-CPU device");
}
phi::Place base_place = options._PD_GetPlace();
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(base_place);
auto dense = paddle::experimental::ones(
size._PD_ToPaddleIntArray(),
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
phi::CPUPlace());
return dense.copy_to(pinned_place, /*blocking=*/true);
}
return paddle::experimental::ones(
size._PD_ToPaddleIntArray(),
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
options._PD_GetPlace());
}
inline at::Tensor ones_symint(c10::SymIntArrayRef size,
::std::optional<at::ScalarType> dtype,
::std::optional<at::Layout> layout,
::std::optional<at::Device> device,
::std::optional<bool> pin_memory) {
PD_CHECK(!layout.has_value(), "`layout` is not supported now.");
auto options =
at::TensorOptions()
.dtype(dtype.value_or(c10::get_default_dtype_as_scalartype()))
.device(device.value_or(at::kCPU))
.pinned_memory(pin_memory);
return ones_symint(size, options);
}
} // namespace at
@@ -0,0 +1,37 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
namespace at {
inline at::Tensor permute(const at::Tensor& self, at::IntArrayRef dims) {
std::vector<int> perm(dims.size());
for (size_t i = 0; i < dims.size(); i++) {
perm[i] = static_cast<int>(dims[i]);
}
return paddle::experimental::transpose(self._PD_GetInner(), perm);
}
} // namespace at
namespace at {
inline at::Tensor Tensor::permute(at::IntArrayRef dims) const {
return at::permute(*this, dims);
}
} // namespace at
@@ -0,0 +1,37 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
namespace at {
inline at::Tensor reciprocal(const at::Tensor& self) {
return Tensor(paddle::experimental::reciprocal(self._PD_GetInner()));
}
} // namespace at
namespace at {
inline at::Tensor Tensor::reciprocal() const { return at::reciprocal(*this); }
inline at::Tensor& Tensor::reciprocal_() const {
PaddleTensor& inner = const_cast<PaddleTensor&>(tensor_);
paddle::experimental::reciprocal_(inner);
return const_cast<at::Tensor&>(*this);
}
} // namespace at
@@ -0,0 +1,52 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/core/Device.h>
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
#include <c10/cuda/CUDAStream.h>
#endif
namespace at {
inline void Tensor::record_stream(at::Stream s) const {
auto dense_tensor =
std::dynamic_pointer_cast<phi::DenseTensor>(tensor_.impl());
PD_CHECK(dense_tensor != nullptr,
"record_stream only supports DenseTensor, but got a non-dense "
"tensor implementation.");
PD_CHECK(dense_tensor->place().GetType() != phi::AllocationType::CPU,
"record_stream is not supported for CPU tensors.");
#if (defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)) && \
!defined(PADDLE_WITH_CUSTOM_DEVICE)
paddle::memory::RecordStream(
dense_tensor->Holder(), reinterpret_cast<gpuStream_t>(s.native_handle()));
#elif defined(PADDLE_WITH_XPU)
paddle::memory::RecordStream(dense_tensor->Holder(),
reinterpret_cast<XPUStream>(s.native_handle()));
#elif defined(PADDLE_WITH_CUSTOM_DEVICE)
paddle::memory::RecordStream(
dense_tensor->Holder(),
reinterpret_cast<phi::stream::stream_t>(s.native_handle()));
#else
(void)s;
(void)dense_tensor;
PD_THROW(
"record_stream is not supported: no GPU/XPU/Custom device enabled "
"in this build.");
#endif
}
} // namespace at
@@ -0,0 +1,26 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
namespace at {
// Member function: Tensor::rename
inline Tensor Tensor::rename(::std::optional<at::DimnameList>) const {
return *this;
}
} // namespace at
@@ -0,0 +1,44 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/core/TensorOptions.h>
#include <optional>
#include <string_view>
#include "paddle/phi/api/include/api.h"
namespace at {
inline at::Tensor reshape(const at::Tensor& self, at::IntArrayRef shape) {
return paddle::experimental::reshape(self._PD_GetInner(),
shape._PD_ToPaddleIntArray());
}
inline at::Tensor reshape_symint(const at::Tensor& self,
c10::SymIntArrayRef shape) {
return paddle::experimental::reshape(self._PD_GetInner(),
shape._PD_ToPaddleIntArray());
}
} // namespace at
namespace at {
inline at::Tensor Tensor::reshape(at::IntArrayRef shape) const {
return at::reshape(*this, shape);
}
} // namespace at
@@ -0,0 +1,119 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/core/TensorOptions.h>
#include <limits>
#include <optional>
#include <string_view>
#include <vector>
#include "paddle/phi/api/include/api.h"
#include "paddle/phi/common/memory_utils.h"
#include "paddle/phi/core/ddim.h"
#include "paddle/phi/core/memory/malloc.h"
namespace at {
namespace detail {
inline int64_t ResizeCheckedNumel(at::IntArrayRef size) {
int64_t numel = 1;
for (const auto dim : size) {
TORCH_CHECK(dim >= 0,
"Trying to create tensor with negative dimension ",
dim,
": ",
size);
if (dim == 0) {
numel = 0;
continue;
}
TORCH_CHECK(numel <= std::numeric_limits<int64_t>::max() / dim,
"resize_ size is too large, possible overflow for size ",
size);
numel *= dim;
}
return numel;
}
inline size_t ResizeCheckedStorageBytes(int64_t numel,
size_t itemsize,
size_t storage_offset_bytes) {
const auto numel_size = static_cast<size_t>(numel);
TORCH_CHECK(
itemsize == 0 || numel_size <= (std::numeric_limits<size_t>::max() -
storage_offset_bytes) /
itemsize,
"resize_ size is too large in bytes");
return storage_offset_bytes + numel_size * itemsize;
}
} // namespace detail
// resize_ - operate on the underlying DenseTensor directly so we preserve
// storage semantics across shrink/grow round-trips. When growth exceeds the
// current capacity, expand the shared storage itself so aliasing views keep
// their storage offset and existing storage contents stay intact.
inline const at::Tensor& Tensor::resize_(
at::IntArrayRef size,
::std::optional<at::MemoryFormat> memory_format) const {
// Keep old compat behavior for memory_format in this split PR.
// TODO(youge325): add real ChannelsLast/ChannelsLast3d restride support
// later.
(void)memory_format;
std::vector<int64_t> dims(size.begin(), size.end());
int64_t new_numel = detail::ResizeCheckedNumel(size);
auto dense_tensor =
std::dynamic_pointer_cast<phi::DenseTensor>(tensor_.impl());
TORCH_CHECK(dense_tensor != nullptr,
"resize_ only supports DenseTensor, but got a non-dense tensor");
TORCH_CHECK(tensor_.defined(),
"resize_ is not allowed on an undefined tensor");
const size_t itemsize = phi::SizeOf(dense_tensor->dtype());
const size_t new_storage_bytes = detail::ResizeCheckedStorageBytes(
new_numel, itemsize, dense_tensor->meta().offset);
const size_t current_storage_bytes =
dense_tensor->Holder() == nullptr ? 0 : dense_tensor->Holder()->size();
if (new_storage_bytes <= current_storage_bytes || new_numel == 0) {
dense_tensor->Resize(dims);
return *this;
}
// Sync through the compat Storage path first so the DenseTensor holder is a
// live StorageHolderView backed by shared StorageImpl.
auto storage = this->storage();
const auto old_holder = dense_tensor->Holder();
TORCH_CHECK(old_holder != nullptr,
"resize_ cannot grow a tensor without allocated storage");
const phi::Place place = old_holder->place();
auto new_holder = paddle::memory::AllocShared(place, new_storage_bytes);
TORCH_CHECK(new_holder != nullptr, "resize_ failed to allocate storage");
const size_t copy_bytes = std::min(old_holder->size(), new_storage_bytes);
if (copy_bytes > 0 && old_holder->ptr() != nullptr &&
old_holder->ptr() != new_holder->ptr()) {
phi::memory_utils::Copy(
place, new_holder->ptr(), place, old_holder->ptr(), copy_bytes);
}
storage.set_data_ptr_noswap(std::move(new_holder));
dense_tensor->Resize(phi::make_ddim(dims));
return *this;
}
} // namespace at
@@ -0,0 +1,80 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
namespace at {
inline at::Tensor select(const at::Tensor& self, int64_t dim, int64_t index) {
// Normalize dim to positive value for error messages
int64_t orig_dim = dim;
if (dim < 0) {
dim += self.dim();
}
// Check dim is valid
if (dim < 0 || dim >= self.dim()) {
PD_CHECK(false,
"select(): index ",
orig_dim,
" out of range for tensor of size ",
self.sizes(),
" at dimension ",
orig_dim);
}
// Handle negative index
int64_t orig_index = index;
if (index < 0) {
index = self.size(dim) + index;
}
// Check index is valid
if (index < 0 || index >= self.size(dim)) {
PD_CHECK(false,
"select(): index ",
orig_index,
" out of range for tensor of size ",
self.sizes(),
" at dimension ",
orig_dim < 0 ? orig_dim + self.dim() : orig_dim);
}
return Tensor(
paddle::experimental::slice(self._PD_GetInner(),
/*axes=*/{static_cast<int>(dim)},
/*starts=*/{index},
/*ends=*/{index + 1},
/*infer_flags=*/{1},
/*decrease_axis=*/{static_cast<int>(dim)}));
}
inline at::Tensor select_symint(const at::Tensor& self,
int64_t dim,
c10::SymInt index) {
return select(self, dim, static_cast<int64_t>(index));
}
} // namespace at
namespace at {
inline at::Tensor Tensor::select(int64_t dim, int64_t index) const {
return at::select(*this, dim, index);
}
inline at::Tensor Tensor::select_symint(int64_t dim, c10::SymInt index) const {
return at::select_symint(*this, dim, index);
}
} // namespace at
@@ -0,0 +1,51 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
namespace at {
inline at::Tensor slice(const at::Tensor& self,
int64_t dim = 0,
::std::optional<int64_t> start = ::std::nullopt,
::std::optional<int64_t> end = ::std::nullopt,
int64_t step = 1) {
// Materialize the compat StorageHolderView before creating the slice so the
// base tensor and its views observe the same shared storage during resize_.
(void)self.storage();
return paddle::experimental::slice(
self._PD_GetInner(),
{dim},
start.has_value() ? IntArrayRef(start.value())._PD_ToPaddleIntArray()
: IntArrayRef()._PD_ToPaddleIntArray(),
end.has_value() ? IntArrayRef(end.value())._PD_ToPaddleIntArray()
: IntArrayRef()._PD_ToPaddleIntArray(),
{1},
{});
}
} // namespace at
namespace at {
inline at::Tensor Tensor::slice(int64_t dim,
::std::optional<int64_t> start,
::std::optional<int64_t> end,
int64_t step) const {
return at::slice(*this, dim, start, end, step);
}
} // namespace at
@@ -0,0 +1,119 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/core/TensorOptions.h>
#include <utils/pinned_place.h>
#include <algorithm>
#include <memory>
#include <optional>
#include "paddle/phi/api/include/api.h"
#include "paddle/phi/api/include/sparse_api.h"
#include "paddle/phi/common/place.h"
#include "paddle/phi/core/sparse_coo_tensor.h"
namespace at::detail {
inline std::vector<int64_t> _PD_infer_sparse_coo_size(
const at::Tensor& indices) {
auto host_indices = indices.cpu().to(at::kLong);
int64_t sparse_dim = host_indices.dim() > 0 ? host_indices.size(0) : 0;
int64_t nnz = host_indices.dim() > 1 ? host_indices.size(1) : 0;
std::vector<int64_t> inferred_size(static_cast<size_t>(sparse_dim), 0);
const int64_t* data = host_indices.const_data_ptr<int64_t>();
for (int64_t dim = 0; dim < sparse_dim; ++dim) {
for (int64_t i = 0; i < nnz; ++i) {
inferred_size[static_cast<size_t>(dim)] = std::max(
inferred_size[static_cast<size_t>(dim)], data[dim * nnz + i] + 1);
}
}
return inferred_size;
}
inline void _PD_set_sparse_coo_coalesced(at::Tensor* tensor,
::std::optional<bool> is_coalesced) {
if (!is_coalesced.has_value()) {
return;
}
auto sparse_tensor = std::dynamic_pointer_cast<phi::SparseCooTensor>(
tensor->_PD_GetInner().impl());
PD_CHECK(sparse_tensor,
"Expected SparseCooTensor result from sparse_coo_tensor.");
sparse_tensor->SetCoalesced(is_coalesced.value());
}
} // namespace at::detail
namespace at {
inline at::Tensor sparse_coo_tensor(
const at::Tensor& indices,
const at::Tensor& values,
at::IntArrayRef size,
at::TensorOptions options = {},
::std::optional<bool> is_coalesced = ::std::nullopt) {
paddle::Tensor idx = indices._PD_GetInner();
paddle::Tensor vals = values._PD_GetInner();
// PyTorch ignores dtype mismatch between values and TensorOptions in
// sparse_coo_tensor; the resulting sparse tensor uses values' original dtype.
// Do not cast or throw here.
if (options.pinned_memory()) {
phi::Place base_place = options._PD_GetPlace();
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(base_place);
idx = idx.copy_to(pinned_place, /*blocking=*/true);
vals = vals.copy_to(pinned_place, /*blocking=*/true);
}
// PyTorch: sparse_coo_tensor(indices, values, size)
// Paddle: sparse_coo_tensor(values, indices, shape)
at::Tensor result = paddle::experimental::sparse::sparse_coo_tensor(
vals, idx, std::vector<int64_t>(size.begin(), size.end()));
detail::_PD_set_sparse_coo_coalesced(&result, is_coalesced);
return result;
}
inline at::Tensor sparse_coo_tensor(const at::Tensor& indices,
const at::Tensor& values,
at::IntArrayRef size,
::std::optional<at::ScalarType> dtype,
::std::optional<at::Layout> layout,
::std::optional<at::Device> device,
::std::optional<bool> pin_memory,
::std::optional<bool> is_coalesced) {
PD_CHECK(!layout.has_value() || layout.value() == c10::kSparse,
"`layout` must be Sparse for sparse_coo_tensor.");
auto options =
at::TensorOptions().dtype(dtype).device(device).pinned_memory(pin_memory);
return sparse_coo_tensor(indices, values, size, options, is_coalesced);
}
inline at::Tensor sparse_coo_tensor(
const at::Tensor& indices,
const at::Tensor& values,
at::TensorOptions options = {},
::std::optional<bool> is_coalesced = ::std::nullopt) {
return sparse_coo_tensor(indices,
values,
detail::_PD_infer_sparse_coo_size(indices),
options,
is_coalesced);
}
} // namespace at
@@ -0,0 +1,137 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/core/TensorOptions.h>
#include <utils/pinned_place.h>
#include <optional>
#include "paddle/phi/api/include/api.h"
#include "paddle/phi/api/include/tensor.h"
#include "paddle/phi/common/place.h"
#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/core/sparse_csr_tensor.h"
namespace at {
inline at::Tensor sparse_csr_tensor(const at::Tensor& crow_indices,
const at::Tensor& col_indices,
const at::Tensor& values,
at::IntArrayRef size,
at::TensorOptions options) {
paddle::Tensor crows = crow_indices._PD_GetInner();
paddle::Tensor cols = col_indices._PD_GetInner();
paddle::Tensor vals = values._PD_GetInner();
// PyTorch ignores dtype mismatch between values and TensorOptions in
// sparse_csr_tensor; the resulting sparse tensor uses values' original dtype.
// Do not cast or throw here.
if (options.pinned_memory()) {
phi::Place base_place = options._PD_GetPlace();
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(base_place);
crows = crows.copy_to(pinned_place, /*blocking=*/true);
cols = cols.copy_to(pinned_place, /*blocking=*/true);
vals = vals.copy_to(pinned_place, /*blocking=*/true);
}
// Get the underlying DenseTensors
auto* dense_crows = dynamic_cast<phi::DenseTensor*>(crows.impl().get());
auto* dense_cols = dynamic_cast<phi::DenseTensor*>(cols.impl().get());
auto* dense_values = dynamic_cast<phi::DenseTensor*>(vals.impl().get());
PD_CHECK(dense_crows != nullptr,
"crow_indices must be a dense tensor for sparse_csr_tensor.");
PD_CHECK(dense_cols != nullptr,
"col_indices must be a dense tensor for sparse_csr_tensor.");
PD_CHECK(dense_values != nullptr,
"values must be a dense tensor for sparse_csr_tensor.");
// Create the SparseCsrTensor
std::shared_ptr<phi::SparseCsrTensor> csr_tensor =
std::make_shared<phi::SparseCsrTensor>(
*dense_crows,
*dense_cols,
*dense_values,
common::make_ddim(std::vector<int64_t>(size.begin(), size.end())));
// Wrap in a Paddle Tensor
paddle::Tensor result;
result.set_impl(csr_tensor);
return result;
}
inline at::Tensor sparse_csr_tensor(const at::Tensor& crow_indices,
const at::Tensor& col_indices,
const at::Tensor& values,
at::IntArrayRef size,
::std::optional<at::ScalarType> dtype,
::std::optional<at::Layout> layout,
::std::optional<at::Device> device,
::std::optional<bool> pin_memory) {
PD_CHECK(!layout.has_value() || layout.value() == c10::kSparseCsr,
"`layout` must be SparseCsr for sparse_csr_tensor.");
auto options =
at::TensorOptions().dtype(dtype).device(device).pinned_memory(pin_memory);
return sparse_csr_tensor(crow_indices, col_indices, values, size, options);
}
inline at::Tensor sparse_csr_tensor(const at::Tensor& crow_indices,
const at::Tensor& col_indices,
const at::Tensor& values,
at::TensorOptions options) {
// Infer size from crow_indices and col_indices:
// nrows = crow_indices.size(0) - 1
// ncols = max(col_indices) + 1
int64_t nrows = crow_indices.size(0) - 1;
int64_t ncols = 0;
if (col_indices.numel() > 0) {
auto* dense_cols = dynamic_cast<phi::DenseTensor*>(
col_indices._PD_GetInner().impl().get());
PD_CHECK(dense_cols != nullptr,
"col_indices must be a dense tensor for sparse_csr_tensor.");
PD_CHECK(
dense_cols->place().GetType() == phi::AllocationType::CPU,
"sparse_csr_tensor without explicit size only supports CPU "
"col_indices for automatic size inference. Please provide the size "
"parameter explicitly for non-CPU tensors.");
int64_t n = dense_cols->numel();
if (dense_cols->dtype() == phi::DataType::INT64) {
const int64_t* data = dense_cols->data<int64_t>();
for (int64_t i = 0; i < n; ++i) {
if (data[i] + 1 > ncols) ncols = data[i] + 1;
}
} else if (dense_cols->dtype() == phi::DataType::INT32) {
const int32_t* data = dense_cols->data<int32_t>();
for (int64_t i = 0; i < n; ++i) {
int64_t val = static_cast<int64_t>(data[i]) + 1;
if (val > ncols) ncols = val;
}
} else {
PD_CHECK(false,
"col_indices must have dtype int32 or int64 for automatic "
"size inference in sparse_csr_tensor.");
}
}
std::vector<int64_t> size_vec = {nrows, ncols};
return sparse_csr_tensor(
crow_indices, col_indices, values, at::IntArrayRef(size_vec), options);
}
} // namespace at
@@ -0,0 +1,93 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
namespace at {
inline std::vector<at::Tensor> split(const at::Tensor& self,
int64_t split_size,
int64_t dim = 0) {
// Calculate number of splits based on split_size
int64_t dim_size = self._PD_GetInner().dims()[dim];
std::vector<int64_t> split_sizes;
for (int64_t i = 0; i < dim_size; i += split_size) {
split_sizes.push_back(std::min(split_size, dim_size - i));
}
auto outputs =
paddle::experimental::split(self._PD_GetInner(), split_sizes, dim);
std::vector<at::Tensor> at_tensors;
at_tensors.reserve(outputs.size());
for (const auto& paddle_tensor : outputs) {
at_tensors.emplace_back(paddle_tensor);
}
return at_tensors;
}
inline std::vector<at::Tensor> split_symint(const at::Tensor& self,
c10::SymInt split_size,
int64_t dim = 0) {
return split(self, static_cast<int64_t>(split_size), dim);
}
inline std::vector<at::Tensor> split(const at::Tensor& self,
at::IntArrayRef split_sizes,
int64_t dim = 0) {
auto outputs = paddle::experimental::split(
self._PD_GetInner(), split_sizes._PD_ToPaddleIntArray(), dim);
std::vector<at::Tensor> at_tensors;
at_tensors.reserve(outputs.size());
for (const auto& paddle_tensor : outputs) {
at_tensors.emplace_back(paddle_tensor);
}
return at_tensors;
}
inline std::vector<at::Tensor> split_symint(const at::Tensor& self,
c10::SymIntArrayRef split_sizes,
int64_t dim = 0) {
return split(
self,
at::IntArrayRef(reinterpret_cast<const int64_t*>(split_sizes.data()),
split_sizes.size()),
dim);
}
} // namespace at
namespace at {
inline std::vector<at::Tensor> Tensor::split(int64_t split_size,
int64_t dim = 0) const {
return at::split(*this, split_size, dim);
}
inline std::vector<at::Tensor> Tensor::split_symint(c10::SymInt split_size,
int64_t dim = 0) const {
return at::split_symint(*this, split_size, dim);
}
inline std::vector<at::Tensor> Tensor::split(at::IntArrayRef split_sizes,
int64_t dim = 0) const {
return at::split(*this, split_sizes, dim);
}
inline std::vector<at::Tensor> Tensor::split_symint(
c10::SymIntArrayRef split_sizes, int64_t dim = 0) const {
return at::split_symint(*this, split_sizes, dim);
}
} // namespace at
@@ -0,0 +1,47 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <ATen/ops/split.h>
namespace at {
inline std::vector<at::Tensor> split_with_sizes(const at::Tensor& self,
at::IntArrayRef split_sizes,
int64_t dim = 0) {
return at::split(self, split_sizes, dim);
}
inline std::vector<at::Tensor> split_with_sizes_symint(
const at::Tensor& self, c10::SymIntArrayRef split_sizes, int64_t dim = 0) {
return at::split_symint(self, split_sizes, dim);
}
} // namespace at
namespace at {
inline std::vector<at::Tensor> Tensor::split_with_sizes(
at::IntArrayRef split_sizes, int64_t dim) const {
return at::split_with_sizes(*this, split_sizes, dim);
}
inline std::vector<at::Tensor> Tensor::split_with_sizes_symint(
c10::SymIntArrayRef split_sizes, int64_t dim) const {
return at::split_with_sizes_symint(*this, split_sizes, dim);
}
} // namespace at
@@ -0,0 +1,66 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
namespace at {
inline at::Tensor squeeze(const at::Tensor& self) {
return paddle::experimental::squeeze(self._PD_GetInner(), {});
}
inline at::Tensor squeeze(const at::Tensor& self, int64_t dim) {
return paddle::experimental::squeeze(self._PD_GetInner(), {dim});
}
inline at::Tensor squeeze(const at::Tensor& self, at::IntArrayRef dim) {
return paddle::experimental::squeeze(self._PD_GetInner(),
dim._PD_ToPaddleIntArray());
}
} // namespace at
namespace at {
inline at::Tensor Tensor::squeeze() const { return at::squeeze(*this); }
inline at::Tensor Tensor::squeeze(int64_t dim) const {
return at::squeeze(*this, dim);
}
inline at::Tensor Tensor::squeeze(at::IntArrayRef dim) const {
return at::squeeze(*this, dim);
}
inline at::Tensor& Tensor::squeeze_() const {
PaddleTensor& inner = const_cast<PaddleTensor&>(tensor_);
paddle::experimental::squeeze_(inner, {});
return const_cast<at::Tensor&>(*this);
}
inline at::Tensor& Tensor::squeeze_(int64_t dim) const {
PaddleTensor& inner = const_cast<PaddleTensor&>(tensor_);
paddle::experimental::squeeze_(inner, {dim});
return const_cast<at::Tensor&>(*this);
}
inline at::Tensor& Tensor::squeeze_(at::IntArrayRef dim) const {
PaddleTensor& inner = const_cast<PaddleTensor&>(tensor_);
paddle::experimental::squeeze_(inner, dim._PD_ToPaddleIntArray());
return const_cast<at::Tensor&>(*this);
}
} // namespace at
@@ -0,0 +1,146 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/core/Scalar.h>
#include <c10/util/ArrayRef.h>
#include <c10/util/OptionalArrayRef.h>
#include <optional>
#include <vector>
#include "paddle/phi/api/include/api.h"
#include "paddle/phi/common/int_array.h"
#include "paddle/phi/common/scalar.h"
namespace at::detail {
// Internal implementation for std (standard deviation = sqrt(variance))
inline Tensor _PD_std_impl(const Tensor& self,
const std::vector<int64_t>& dims_vec,
double correction_value,
bool keepdim) {
// Validate dimensions before processing
int64_t ndim = self.dim();
for (int64_t d : dims_vec) {
int64_t dim_idx = d < 0 ? d + ndim : d;
if (dim_idx < 0 || dim_idx >= ndim) {
PD_CHECK(false,
"Dimension out of range (expected to be in range of [",
-ndim,
", ",
ndim - 1,
"], but got ",
d,
")");
}
}
phi::IntArray dims_int_array(dims_vec);
paddle::Tensor tensor = self._PD_GetInner();
paddle::Tensor mean_tensor;
if (dims_vec.empty()) {
mean_tensor = paddle::experimental::mean(
tensor, phi::IntArray(std::vector<int64_t>{}), true);
} else {
mean_tensor = paddle::experimental::mean(tensor, dims_int_array, true);
}
paddle::Tensor diff = paddle::experimental::subtract(tensor, mean_tensor);
paddle::Tensor diff_squared = paddle::experimental::multiply(diff, diff);
paddle::Tensor sum_squared_diff;
if (dims_vec.empty()) {
sum_squared_diff =
paddle::experimental::sum(diff_squared,
phi::IntArray(std::vector<int64_t>{}),
diff_squared.dtype(),
keepdim);
} else {
sum_squared_diff = paddle::experimental::sum(
diff_squared, dims_int_array, diff_squared.dtype(), keepdim);
}
int64_t n = tensor.numel();
if (!dims_vec.empty()) {
n = 1;
for (int64_t d : dims_vec) {
int64_t dim_idx = d < 0 ? d + tensor.dims().size() : d;
if (dim_idx >= 0 &&
dim_idx < static_cast<int64_t>(tensor.dims().size())) {
n *= tensor.dims()[dim_idx];
}
}
}
double corrected_n = static_cast<double>(n) - correction_value;
if (corrected_n <= 0.0) {
corrected_n = static_cast<double>(n);
}
std::vector<int64_t> result_shape_vec;
for (int64_t i = 0; i < sum_squared_diff.dims().size(); ++i) {
result_shape_vec.push_back(sum_squared_diff.dims()[i]);
}
paddle::Tensor correction_scalar =
paddle::experimental::full(phi::IntArray(result_shape_vec),
phi::Scalar(corrected_n),
sum_squared_diff.dtype(),
sum_squared_diff.place());
paddle::Tensor variance =
paddle::experimental::divide(sum_squared_diff, correction_scalar);
paddle::Tensor result = paddle::experimental::sqrt(variance);
return Tensor(result);
}
} // namespace at::detail
namespace at {
inline Tensor Tensor::std(bool unbiased) const {
std::vector<int64_t> empty_dims;
double correction = unbiased ? 1.0 : 0.0;
return detail::_PD_std_impl(*this, empty_dims, correction, false);
}
inline Tensor Tensor::std(at::OptionalIntArrayRef dim,
bool unbiased,
bool keepdim) const {
double correction = unbiased ? 1.0 : 0.0;
std::vector<int64_t> dims_vec;
if (dim.has_value() && dim.value().size() > 0) {
dims_vec.assign(dim.value().begin(), dim.value().end());
}
return detail::_PD_std_impl(*this, dims_vec, correction, keepdim);
}
inline Tensor Tensor::std(at::OptionalIntArrayRef dim,
const ::std::optional<at::Scalar>& correction,
bool keepdim) const {
double correction_value = 1.0;
if (correction.has_value()) {
const at::Scalar& scalar = correction.value();
correction_value = scalar.to<double>();
}
std::vector<int64_t> dims_vec;
if (dim.has_value() && dim.value().size() > 0) {
dims_vec.assign(dim.value().begin(), dim.value().end());
}
return detail::_PD_std_impl(*this, dims_vec, correction_value, keepdim);
}
} // namespace at
@@ -0,0 +1,110 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/core/TensorOptions.h>
#include <c10/util/OptionalArrayRef.h>
#include <optional>
#include <string_view>
#include "paddle/phi/api/include/api.h"
namespace at {
inline at::Tensor sum(const at::Tensor& self,
::std::optional<at::ScalarType> dtype = ::std::nullopt) {
// Match PyTorch promotion: integer inputs -> int64; others -> keep input
// dtype.
at::ScalarType resolved_dtype;
if (dtype.has_value()) {
resolved_dtype = dtype.value();
} else {
at::ScalarType input_dtype = self.scalar_type();
resolved_dtype = at::isIntegralType(input_dtype, /*includeBool=*/true)
? at::kLong
: input_dtype;
}
return paddle::experimental::sum(
self._PD_GetInner(),
{},
compat::_PD_AtenScalarTypeToPhiDataType(resolved_dtype),
/*keepdim=*/false);
}
inline at::Tensor sum(const at::Tensor& self,
at::OptionalIntArrayRef dim,
bool keepdim = false,
::std::optional<at::ScalarType> dtype = ::std::nullopt) {
// Match PyTorch promotion: integer inputs -> int64; others -> keep input
// dtype.
at::ScalarType resolved_dtype;
if (dtype.has_value()) {
resolved_dtype = dtype.value();
} else {
at::ScalarType input_dtype = self.scalar_type();
resolved_dtype = at::isIntegralType(input_dtype, /*includeBool=*/true)
? at::kLong
: input_dtype;
}
return paddle::experimental::sum(
self._PD_GetInner(),
dim.has_value() ? dim.value()._PD_ToPaddleIntArray()
: paddle::experimental::IntArray(),
compat::_PD_AtenScalarTypeToPhiDataType(resolved_dtype),
keepdim);
}
inline at::Tensor& sum_out(
at::Tensor&
out, // NOLINT: intentional non-const reference for output parameter
const at::Tensor& self,
at::OptionalIntArrayRef dim,
bool keepdim = false,
::std::optional<at::ScalarType> dtype = ::std::nullopt) {
auto res = sum(self, dim, keepdim, dtype);
paddle::experimental::assign_out_(res._PD_GetInner(), out._PD_GetInner());
return out;
}
inline at::Tensor& sum_out(
at::Tensor&
out, // NOLINT: intentional non-const reference for output parameter
const at::Tensor& self,
::std::optional<at::ScalarType> dtype = ::std::nullopt) {
auto res = sum(self, dtype);
paddle::experimental::assign_out_(res._PD_GetInner(), out._PD_GetInner());
return out;
}
} // namespace at
namespace at {
inline at::Tensor Tensor::sum(::std::optional<at::ScalarType> dtype) const {
return at::sum(*this, dtype);
}
inline at::Tensor Tensor::sum(at::OptionalIntArrayRef dim,
bool keepdim,
::std::optional<at::ScalarType> dtype) const {
return at::sum(*this, dim, keepdim, dtype);
}
} // namespace at
@@ -0,0 +1,35 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
namespace at {
inline at::Tensor t(const Tensor& self) {
return self.transpose(0, self.dim() < 2 ? 0 : 1);
}
} // namespace at
namespace at {
inline at::Tensor Tensor::t() const { return at::t(*this); }
inline at::Tensor& Tensor::t_() const {
return transpose_(0, dim() < 2 ? 0 : 1);
}
} // namespace at
@@ -0,0 +1,47 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/core/ScalarType.h>
#include "paddle/common/macros.h"
namespace at {
#define TENSOR(T, S) \
PADDLE_API Tensor tensor(ArrayRef<T> values, const TensorOptions& options); \
inline Tensor tensor(std::initializer_list<T> values, \
const TensorOptions& options) { \
return at::tensor(ArrayRef<T>(values), options); \
} \
inline Tensor tensor(T value, const TensorOptions& options) { \
return at::tensor(ArrayRef<T>(value), options); \
} \
inline Tensor tensor(ArrayRef<T> values) { \
return at::tensor(std::move(values), at::dtype(k##S)); \
} \
inline Tensor tensor(std::initializer_list<T> values) { \
return at::tensor(ArrayRef<T>(values)); \
} \
inline Tensor tensor(T value) { return at::tensor(ArrayRef<T>(value)); }
AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, TENSOR)
AT_FORALL_COMPLEX_TYPES(TENSOR)
#undef TENSOR
} // namespace at
@@ -0,0 +1,207 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// The file has been adapted from pytorch project
// Licensed under BSD-style license -
// https://github.com/pytorch/pytorch/blob/main/LICENSE
#pragma once
#include <ATen/core/Tensor.h>
namespace at {
inline std::vector<at::Tensor> tensor_split(const at::Tensor& self,
int64_t sections,
int64_t dim = 0) {
// Follow PyTorch's tensor_split_sections_symint implementation
PD_CHECK(self._PD_GetInner().dims().size() > 0,
"tensor_split expected at least a 1-dimensional tensor, but got a "
"tensor with ",
self._PD_GetInner().dims().size(),
" dims");
PD_CHECK(
sections > 0, "number of sections must be larger than 0, got ", sections);
int64_t dim_size = self._PD_GetInner().dims()[dim];
// Calculate split sizes: first (dim_size % sections) chunks get size
// (dim_size / sections + 1), remaining chunks get size (dim_size / sections)
auto min_split_size = dim_size / sections;
auto num_splits_one_extra = dim_size % sections;
std::vector<int64_t> split_sizes;
split_sizes.reserve(sections);
for (int64_t split_idx = 0; split_idx < sections; ++split_idx) {
auto split_size = (split_idx < num_splits_one_extra) ? (min_split_size + 1)
: min_split_size;
split_sizes.push_back(split_size);
}
// Use split with calculated sizes
auto outputs =
paddle::experimental::split(self._PD_GetInner(), split_sizes, dim);
std::vector<at::Tensor> at_tensors;
at_tensors.reserve(outputs.size());
for (const auto& paddle_tensor : outputs) {
at_tensors.emplace_back(paddle_tensor);
}
return at_tensors;
}
inline std::vector<at::Tensor> tensor_split_symint(const at::Tensor& self,
c10::SymInt sections,
int64_t dim = 0) {
return tensor_split(self, static_cast<int64_t>(sections), dim);
}
inline std::vector<at::Tensor> tensor_split(const at::Tensor& self,
at::IntArrayRef indices,
int64_t dim = 0) {
// Follow PyTorch's _tensor_split_indices implementation
// indices are split positions, not sizes
PD_CHECK(self._PD_GetInner().dims().size() > 0,
"tensor_split expected at least a 1-dimensional tensor, but got a "
"tensor with ",
self._PD_GetInner().dims().size(),
" dims");
int64_t num_indices = indices.size();
int64_t dim_size = self._PD_GetInner().dims()[dim];
// Convert indices (positions) to sizes
std::vector<int64_t> split_sizes;
split_sizes.reserve(num_indices + 1);
int64_t start_idx = 0;
for (int64_t i = 0; i < num_indices; ++i) {
int64_t end_idx = indices[i];
// Handle negative indices
if (end_idx < 0) {
end_idx += dim_size;
}
// Clamp to valid range
end_idx = std::max(start_idx, std::min(end_idx, dim_size));
split_sizes.push_back(end_idx - start_idx);
start_idx = end_idx;
}
// Add the last segment
split_sizes.push_back(dim_size - start_idx);
// Use split with calculated sizes
auto outputs =
paddle::experimental::split(self._PD_GetInner(), split_sizes, dim);
std::vector<at::Tensor> at_tensors;
at_tensors.reserve(outputs.size());
for (const auto& paddle_tensor : outputs) {
at_tensors.emplace_back(paddle_tensor);
}
return at_tensors;
}
inline std::vector<at::Tensor> tensor_split_symint(const at::Tensor& self,
c10::SymIntArrayRef indices,
int64_t dim = 0) {
return tensor_split(
self,
at::IntArrayRef(reinterpret_cast<const int64_t*>(indices.data()),
indices.size()),
dim);
}
inline std::vector<at::Tensor> tensor_split(
const at::Tensor& self,
const at::Tensor& tensor_indices_or_sections,
int64_t dim = 0) {
// Follow PyTorch's validation and implementation
PD_CHECK(self._PD_GetInner().dims().size() > 0,
"tensor_split expected at least a 1-dimensional tensor, but got a "
"tensor with ",
self._PD_GetInner().dims().size(),
" dims");
auto split_device = tensor_indices_or_sections.device();
PD_CHECK(split_device.is_cpu(),
"tensor_split expected tensor_indices_or_sections to be on cpu, but "
"it's on ",
split_device);
auto split_dtype = tensor_indices_or_sections.scalar_type();
PD_CHECK(split_dtype == at::kLong,
"tensor_split expected tensor_indices_or_sections to have dtype of "
"long, but got ",
split_dtype);
auto split_dim = tensor_indices_or_sections.dim();
PD_CHECK(split_dim == 1 || split_dim == 0,
"tensor_split expected tensor_indices_or_sections to be a "
"zero-dimensional or one-dimensional tensor, but got a tensor with ",
split_dim,
" dims");
if (split_dim == 0) {
// 0-dimensional tensor: treat as sections
int64_t sections = tensor_indices_or_sections.item<int64_t>();
return tensor_split(self, sections, dim);
} else {
// 1-dimensional tensor: treat as indices
// Need to handle non-contiguous tensors properly
const PaddleTensor& paddle_tensor =
tensor_indices_or_sections._PD_GetInner();
const int64_t* indices_data = paddle_tensor.data<int64_t>();
auto stride = tensor_indices_or_sections.stride(0);
auto numel = tensor_indices_or_sections.numel();
std::vector<int64_t> indices(numel);
for (int64_t offset = 0; offset < numel; ++offset) {
// indices tensor could be non-contiguous
indices[offset] = *(indices_data + offset * stride);
}
return tensor_split(self, at::IntArrayRef(indices), dim);
}
}
} // namespace at
namespace at {
inline std::vector<at::Tensor> Tensor::tensor_split(int64_t sections,
int64_t dim = 0) const {
return at::tensor_split(*this, sections, dim);
}
inline std::vector<at::Tensor> Tensor::tensor_split_symint(
c10::SymInt sections, int64_t dim = 0) const {
return at::tensor_split_symint(*this, sections, dim);
}
inline std::vector<at::Tensor> Tensor::tensor_split(at::IntArrayRef indices,
int64_t dim = 0) const {
return at::tensor_split(*this, indices, dim);
}
inline std::vector<at::Tensor> Tensor::tensor_split_symint(
c10::SymIntArrayRef indices, int64_t dim = 0) const {
return at::tensor_split_symint(*this, indices, dim);
}
inline std::vector<at::Tensor> Tensor::tensor_split(
const at::Tensor& tensor_indices_or_sections, int64_t dim = 0) const {
return at::tensor_split(*this, tensor_indices_or_sections, dim);
}
} // namespace at
+133
View File
@@ -0,0 +1,133 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/core/Device.h>
#include <c10/core/ScalarType.h>
#include <c10/core/TensorOptions.h>
#include <utils/scalar_type_conversion.h>
#include "paddle/phi/api/include/api.h"
#include "paddle/phi/common/place.h"
namespace at {
// Overload 1: to(TensorOptions, non_blocking, copy, memory_format)
inline at::Tensor Tensor::to(
at::TensorOptions options,
bool non_blocking,
bool copy,
::std::optional<at::MemoryFormat> memory_format) const {
// Handle device transfer
PaddleTensor result = tensor_;
bool materialized_copy = false;
if (options.has_device()) {
const c10::Device& dev = options.device();
phi::Place place;
switch (dev.type()) {
case c10::DeviceType::CPU:
case c10::DeviceType::CUDA:
case c10::DeviceType::XPU:
case c10::DeviceType::IPU:
case c10::DeviceType::CUSTOM:
place = dev._PD_GetInner();
break;
default:
PD_THROW("Unsupported device type: ", dev.type());
break;
}
if (place != tensor_.place()) {
result = result.copy_to(place, /*blocking=*/!non_blocking);
materialized_copy = true;
}
}
// Handle dtype cast
if (options.has_dtype()) {
auto target_dtype =
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype());
if (target_dtype != result.dtype()) {
result = paddle::experimental::cast(result, target_dtype);
materialized_copy = true;
}
}
if (copy && !materialized_copy) {
result = paddle::experimental::assign(result);
}
return at::Tensor(result);
}
// Overload 2: to(optional<ScalarType>, optional<Layout>, optional<Device>,
// optional<bool> pin_memory, non_blocking, copy, memory_format)
inline at::Tensor Tensor::to(
::std::optional<at::ScalarType> dtype,
::std::optional<at::Layout> layout,
::std::optional<at::Device> device,
::std::optional<bool> pin_memory,
bool non_blocking,
bool copy,
::std::optional<at::MemoryFormat> memory_format) const {
at::TensorOptions options;
if (dtype.has_value()) {
options = options.dtype(dtype.value());
}
if (device.has_value()) {
options = options.device(device.value());
}
if (pin_memory.has_value() && pin_memory.value()) {
options = options.pinned_memory(true);
}
return to(options, non_blocking, copy, memory_format);
}
// Overload 3: to(Device, ScalarType, non_blocking, copy, memory_format)
inline at::Tensor Tensor::to(
at::Device device,
at::ScalarType dtype,
bool non_blocking,
bool copy,
::std::optional<at::MemoryFormat> memory_format) const {
at::TensorOptions options = at::TensorOptions().device(device).dtype(dtype);
return to(options, non_blocking, copy, memory_format);
}
// Overload 4: to(ScalarType, non_blocking, copy, memory_format)
inline at::Tensor Tensor::to(
at::ScalarType dtype,
bool non_blocking,
bool copy,
::std::optional<at::MemoryFormat> memory_format) const {
auto target_dtype = compat::_PD_AtenScalarTypeToPhiDataType(dtype);
if (!copy && target_dtype == tensor_.dtype()) {
return *this;
}
return at::Tensor(paddle::experimental::cast(tensor_, target_dtype));
}
// Overload 5: to(const Tensor& other, non_blocking, copy, memory_format)
inline at::Tensor Tensor::to(
const at::Tensor& other,
bool non_blocking,
bool copy,
::std::optional<at::MemoryFormat> memory_format) const {
at::TensorOptions options =
at::TensorOptions().device(other.device()).dtype(other.scalar_type());
return to(options, non_blocking, copy, memory_format);
}
} // namespace at
@@ -0,0 +1,88 @@
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
#include <c10/core/TensorOptions.h>
#include <limits>
#include <optional>
#include <string_view>
#include <vector>
#include "paddle/phi/api/include/api.h"
namespace at::detail {
inline int _PD_normalize_transpose_dim(int64_t dim,
int64_t ndim,
const char* name) {
int64_t normalized = dim;
if (normalized < 0) {
normalized += ndim;
}
PD_CHECK(normalized >= 0 && normalized < ndim, name, " out of range");
PD_CHECK(normalized <= static_cast<int64_t>(std::numeric_limits<int>::max()),
name,
" out of int range");
return static_cast<int>(normalized);
}
inline std::vector<int> _PD_make_transpose_perm(int64_t ndim, int d0, int d1) {
PD_CHECK(ndim <= static_cast<int64_t>(std::numeric_limits<int>::max()),
"tensor rank out of int range");
std::vector<int> perm(static_cast<size_t>(ndim));
for (int64_t i = 0; i < ndim; ++i) {
perm[static_cast<size_t>(i)] = static_cast<int>(i);
}
std::swap(perm[d0], perm[d1]);
return perm;
}
} // namespace at::detail
namespace at {
inline at::Tensor transpose(const at::Tensor& self,
int64_t dim0,
int64_t dim1) {
int64_t ndim = self.dim();
int d0 = at::detail::_PD_normalize_transpose_dim(dim0, ndim, "dim0");
int d1 = at::detail::_PD_normalize_transpose_dim(dim1, ndim, "dim1");
auto perm = at::detail::_PD_make_transpose_perm(ndim, d0, d1);
return paddle::experimental::transpose(self._PD_GetInner(), perm);
}
} // namespace at
namespace at {
inline at::Tensor Tensor::transpose(int64_t dim0, int64_t dim1) const {
return at::transpose(*this, dim0, dim1);
}
inline at::Tensor& Tensor::transpose_(int64_t dim0, int64_t dim1) const {
int64_t ndim = this->dim();
int d0 = at::detail::_PD_normalize_transpose_dim(dim0, ndim, "dim0");
int d1 = at::detail::_PD_normalize_transpose_dim(dim1, ndim, "dim1");
auto perm = at::detail::_PD_make_transpose_perm(ndim, d0, d1);
PaddleTensor& inner = const_cast<PaddleTensor&>(tensor_);
paddle::experimental::transpose_(inner, perm);
return const_cast<at::Tensor&>(*this);
}
} // namespace at
@@ -0,0 +1,63 @@
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <ATen/core/Tensor.h>
namespace at {
inline at::Tensor unflatten(const at::Tensor& self,
int64_t dim,
at::IntArrayRef sizes) {
// Compute the new shape by replacing the dimension at 'dim' with 'sizes'
int64_t ndim = self._PD_GetInner().dims().size();
int64_t actual_dim = dim < 0 ? dim + ndim : dim;
std::vector<int64_t> new_shape;
for (int64_t i = 0; i < ndim; ++i) {
if (i == actual_dim) {
for (auto s : sizes) {
new_shape.push_back(s);
}
} else {
new_shape.push_back(self._PD_GetInner().dims()[i]);
}
}
return Tensor(paddle::experimental::reshape(self._PD_GetInner(), new_shape));
}
inline at::Tensor unflatten_symint(const at::Tensor& self,
int64_t dim,
c10::SymIntArrayRef sizes) {
return unflatten(
self,
dim,
at::IntArrayRef(reinterpret_cast<const int64_t*>(sizes.data()),
sizes.size()));
}
} // namespace at
namespace at {
inline at::Tensor Tensor::unflatten(int64_t dim, at::IntArrayRef sizes) const {
return at::unflatten(*this, dim, sizes);
}
inline at::Tensor Tensor::unflatten_symint(int64_t dim,
c10::SymIntArrayRef sizes) const {
return at::unflatten_symint(*this, dim, sizes);
}
} // namespace at

Some files were not shown because too many files have changed in this diff Show More