Add files using upload-large-folder tool
Browse filesThis view is limited to 50 files because it contains too many changes. See raw diff
- lib/python3.12/site-packages/torch/include/ATen/AccumulateType.h +173 -0
- lib/python3.12/site-packages/torch/include/ATen/CPUFunctions.h +29 -0
- lib/python3.12/site-packages/torch/include/ATen/CUDAFunctions_inl.h +625 -0
- lib/python3.12/site-packages/torch/include/ATen/CachedTensorUtils.h +24 -0
- lib/python3.12/site-packages/torch/include/ATen/CompositeExplicitAutogradFunctions_inl.h +560 -0
- lib/python3.12/site-packages/torch/include/ATen/CompositeImplicitAutogradFunctions_inl.h +501 -0
- lib/python3.12/site-packages/torch/include/ATen/Dispatch.h +808 -0
- lib/python3.12/site-packages/torch/include/ATen/ExpandBase.h +30 -0
- lib/python3.12/site-packages/torch/include/ATen/Generator.h +2 -0
- lib/python3.12/site-packages/torch/include/ATen/LegacyBatchedFallback.h +25 -0
- lib/python3.12/site-packages/torch/include/ATen/MatrixRef.h +108 -0
- lib/python3.12/site-packages/torch/include/ATen/MemoryOverlap.h +42 -0
- lib/python3.12/site-packages/torch/include/ATen/MetaFunctions.h +29 -0
- lib/python3.12/site-packages/torch/include/ATen/NamedTensorUtils.h +214 -0
- lib/python3.12/site-packages/torch/include/ATen/OpMathType.h +69 -0
- lib/python3.12/site-packages/torch/include/ATen/Operators.h +1389 -0
- lib/python3.12/site-packages/torch/include/ATen/PTThreadPool.h +17 -0
- lib/python3.12/site-packages/torch/include/ATen/PadNd.h +14 -0
- lib/python3.12/site-packages/torch/include/ATen/PythonTorchFunctionTLS.h +36 -0
- lib/python3.12/site-packages/torch/include/ATen/ScalarOps.h +53 -0
- lib/python3.12/site-packages/torch/include/ATen/ScalarType.h +4 -0
- lib/python3.12/site-packages/torch/include/ATen/SmallVector.h +2 -0
- lib/python3.12/site-packages/torch/include/ATen/SparseCsrTensorImpl.h +206 -0
- lib/python3.12/site-packages/torch/include/ATen/SparseTensorImpl.h +421 -0
- lib/python3.12/site-packages/torch/include/ATen/Tensor.h +3 -0
- lib/python3.12/site-packages/torch/include/ATen/TensorGeometry.h +144 -0
- lib/python3.12/site-packages/torch/include/ATen/TensorMeta.h +137 -0
- lib/python3.12/site-packages/torch/include/ATen/TensorNames.h +75 -0
- lib/python3.12/site-packages/torch/include/ATen/TensorSubclassLikeUtils.h +88 -0
- lib/python3.12/site-packages/torch/include/ATen/TensorUtils.h +190 -0
- lib/python3.12/site-packages/torch/include/ATen/WrapDimUtilsMulti.h +44 -0
- lib/python3.12/site-packages/torch/include/ATen/autocast_mode.h +911 -0
- lib/python3.12/site-packages/torch/include/ATen/ceil_div.h +24 -0
- lib/python3.12/site-packages/torch/include/ATen/core/ATenGeneral.h +3 -0
- lib/python3.12/site-packages/torch/include/ATen/core/ATenOpList.h +13 -0
- lib/python3.12/site-packages/torch/include/ATen/core/ATen_fwd.h +46 -0
- lib/python3.12/site-packages/torch/include/ATen/core/ATen_pch.h +165 -0
- lib/python3.12/site-packages/torch/include/ATen/core/Array.h +48 -0
- lib/python3.12/site-packages/torch/include/ATen/core/Backtrace.h +2 -0
- lib/python3.12/site-packages/torch/include/ATen/core/CachingHostAllocator.h +465 -0
- lib/python3.12/site-packages/torch/include/ATen/core/CheckMemoryFormat.h +24 -0
- lib/python3.12/site-packages/torch/include/ATen/core/DeprecatedTypeProperties.h +139 -0
- lib/python3.12/site-packages/torch/include/ATen/core/DeprecatedTypePropertiesRegistry.h +33 -0
- lib/python3.12/site-packages/torch/include/ATen/core/Dict.h +399 -0
- lib/python3.12/site-packages/torch/include/ATen/core/Dict_inl.h +209 -0
- lib/python3.12/site-packages/torch/include/ATen/core/DimVector.h +13 -0
- lib/python3.12/site-packages/torch/include/ATen/core/Dimname.h +48 -0
- lib/python3.12/site-packages/torch/include/ATen/core/DistributionsHelper.h +329 -0
- lib/python3.12/site-packages/torch/include/ATen/core/Formatting.h +25 -0
- lib/python3.12/site-packages/torch/include/ATen/core/Generator.h +191 -0
lib/python3.12/site-packages/torch/include/ATen/AccumulateType.h
ADDED
|
@@ -0,0 +1,173 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
#include <ATen/Config.h>
|
| 3 |
+
#include <c10/core/DeviceType.h>
|
| 4 |
+
#include <c10/core/ScalarType.h>
|
| 5 |
+
#include <c10/util/BFloat16.h>
|
| 6 |
+
#include <c10/util/Float8_e4m3fn.h>
|
| 7 |
+
#include <c10/util/Float8_e4m3fnuz.h>
|
| 8 |
+
#include <c10/util/Float8_e5m2.h>
|
| 9 |
+
#include <c10/util/Float8_e5m2fnuz.h>
|
| 10 |
+
#include <c10/util/Half.h>
|
| 11 |
+
|
| 12 |
+
// Defines the accumulation type for a scalar type.
|
| 13 |
+
// Example:
|
| 14 |
+
// using accscalar_t = acc_type<scalar_t, /*is_cuda*/true>;
|
| 15 |
+
//
|
| 16 |
+
// Accumulation types are an important concept in numeric computing
|
| 17 |
+
// because you frequently want to perform intermediate computations
|
| 18 |
+
// at a higher precision than the input and output precision, to avoid
|
| 19 |
+
// compounding internal rounding errors. Accumulation is the most
|
| 20 |
+
// well-known intermediate computation (it is of great importance for
|
| 21 |
+
// sum reduction and matrix multiply, for example), but in PyTorch
|
| 22 |
+
// acc_type ends up getting used for all sorts of other intermediate
|
| 23 |
+
// computations, so it perhaps would be more accurately (ahem) called an
|
| 24 |
+
// "accurate" type. acc_type is especially important for reduced
|
| 25 |
+
// precision operations like float16 and bfloat16, where relatively
|
| 26 |
+
// benign looking inputs can easily end up overflowing/underflowing.
|
| 27 |
+
//
|
| 28 |
+
// acc_type is parametrized by whether or not you are running on CUDA
|
| 29 |
+
// or not, because on CUDA double precision operations are expensive
|
| 30 |
+
// and so by default, we don't actually want to use double as an
|
| 31 |
+
// acc_type on CUDA. A lot of things are typed out below, but
|
| 32 |
+
// basically, the table is generated by a few rules:
|
| 33 |
+
//
|
| 34 |
+
// If bool:
|
| 35 |
+
// Use 'bool' as acc_type.
|
| 36 |
+
// If floating point:
|
| 37 |
+
// If CUDA, use 'float' as acc_type (unless scalar_t is double),
|
| 38 |
+
// otherwise (CPU) use 'double'
|
| 39 |
+
// If integral:
|
| 40 |
+
// Use 'int64_t' as acc_type
|
| 41 |
+
//
|
| 42 |
+
// You're not forced to use this template; if you happen to know
|
| 43 |
+
// something specific about your use case, you can specify your own
|
| 44 |
+
// desired behavior. This template, however, will give you a reasonable
|
| 45 |
+
// default that will work for all dtypes supported in PyTorch.
|
| 46 |
+
|
| 47 |
+
#if defined(__CUDACC__)
|
| 48 |
+
#include <cuda.h>
|
| 49 |
+
#include <cuda_fp16.h>
|
| 50 |
+
#elif defined(__HIPCC__)
|
| 51 |
+
#include <hip/hip_fp16.h>
|
| 52 |
+
#include <hip/hip_runtime.h>
|
| 53 |
+
#endif
|
| 54 |
+
|
| 55 |
+
namespace at {
|
| 56 |
+
|
| 57 |
+
template <typename T, c10::DeviceType D>
|
| 58 |
+
struct AccumulateTypeDevice {};
|
| 59 |
+
|
| 60 |
+
template <typename T, bool>
|
| 61 |
+
struct AccumulateType {};
|
| 62 |
+
|
| 63 |
+
template <typename T>
|
| 64 |
+
struct AccumulateType<T, false> {
|
| 65 |
+
using type = typename AccumulateTypeDevice<T, c10::DeviceType::CPU>::type;
|
| 66 |
+
};
|
| 67 |
+
|
| 68 |
+
template <typename T>
|
| 69 |
+
struct AccumulateType<T, true> {
|
| 70 |
+
using type = typename AccumulateTypeDevice<T, c10::DeviceType::CUDA>::type;
|
| 71 |
+
};
|
| 72 |
+
|
| 73 |
+
template <typename T, c10::DeviceType device>
|
| 74 |
+
using acc_type_device = typename AccumulateTypeDevice<T, device>::type;
|
| 75 |
+
|
| 76 |
+
template <typename T, bool is_cuda>
|
| 77 |
+
using acc_type = typename AccumulateType<T, is_cuda>::type;
|
| 78 |
+
|
| 79 |
+
#define ACC_TYPE(t, acc_t, device_type) \
|
| 80 |
+
template <> \
|
| 81 |
+
struct AccumulateTypeDevice<t, device_type> { \
|
| 82 |
+
using type = acc_t; \
|
| 83 |
+
};
|
| 84 |
+
#define MPS_ACC_TYPE(t, acc_t) ACC_TYPE(t, acc_t, c10::DeviceType::MPS)
|
| 85 |
+
#define XPU_ACC_TYPE(t, acc_t) ACC_TYPE(t, acc_t, c10::DeviceType::XPU)
|
| 86 |
+
#define CUDA_ACC_TYPE(t, acc_t) ACC_TYPE(t, acc_t, c10::DeviceType::CUDA)
|
| 87 |
+
#define CPU_ACC_TYPE(t, acc_t) ACC_TYPE(t, acc_t, c10::DeviceType::CPU)
|
| 88 |
+
|
| 89 |
+
MPS_ACC_TYPE(BFloat16, float)
|
| 90 |
+
MPS_ACC_TYPE(Half, float)
|
| 91 |
+
MPS_ACC_TYPE(Float8_e5m2, float)
|
| 92 |
+
MPS_ACC_TYPE(Float8_e4m3fn, float)
|
| 93 |
+
MPS_ACC_TYPE(Float8_e5m2fnuz, float)
|
| 94 |
+
MPS_ACC_TYPE(Float8_e4m3fnuz, float)
|
| 95 |
+
MPS_ACC_TYPE(float, float)
|
| 96 |
+
MPS_ACC_TYPE(double, float)
|
| 97 |
+
MPS_ACC_TYPE(int8_t, int64_t)
|
| 98 |
+
MPS_ACC_TYPE(uint8_t, int64_t)
|
| 99 |
+
MPS_ACC_TYPE(char, int64_t)
|
| 100 |
+
MPS_ACC_TYPE(int16_t, int64_t)
|
| 101 |
+
MPS_ACC_TYPE(int32_t, int64_t)
|
| 102 |
+
MPS_ACC_TYPE(int64_t, int64_t)
|
| 103 |
+
MPS_ACC_TYPE(bool, bool)
|
| 104 |
+
MPS_ACC_TYPE(c10::complex<Half>, c10::complex<float>)
|
| 105 |
+
MPS_ACC_TYPE(c10::complex<float>, c10::complex<float>)
|
| 106 |
+
MPS_ACC_TYPE(c10::complex<double>, c10::complex<float>)
|
| 107 |
+
|
| 108 |
+
XPU_ACC_TYPE(BFloat16, float)
|
| 109 |
+
XPU_ACC_TYPE(Half, float)
|
| 110 |
+
XPU_ACC_TYPE(Float8_e5m2, float)
|
| 111 |
+
XPU_ACC_TYPE(Float8_e4m3fn, float)
|
| 112 |
+
XPU_ACC_TYPE(Float8_e5m2fnuz, float)
|
| 113 |
+
XPU_ACC_TYPE(Float8_e4m3fnuz, float)
|
| 114 |
+
XPU_ACC_TYPE(float, float)
|
| 115 |
+
XPU_ACC_TYPE(double, double)
|
| 116 |
+
XPU_ACC_TYPE(int8_t, int64_t)
|
| 117 |
+
XPU_ACC_TYPE(uint8_t, int64_t)
|
| 118 |
+
XPU_ACC_TYPE(char, int64_t)
|
| 119 |
+
XPU_ACC_TYPE(int16_t, int64_t)
|
| 120 |
+
XPU_ACC_TYPE(int32_t, int64_t)
|
| 121 |
+
XPU_ACC_TYPE(int64_t, int64_t)
|
| 122 |
+
XPU_ACC_TYPE(bool, bool)
|
| 123 |
+
XPU_ACC_TYPE(c10::complex<Half>, c10::complex<float>)
|
| 124 |
+
XPU_ACC_TYPE(c10::complex<float>, c10::complex<float>)
|
| 125 |
+
XPU_ACC_TYPE(c10::complex<double>, c10::complex<double>)
|
| 126 |
+
|
| 127 |
+
#if defined(__CUDACC__) || defined(__HIPCC__)
|
| 128 |
+
CUDA_ACC_TYPE(half, float)
|
| 129 |
+
#endif
|
| 130 |
+
CUDA_ACC_TYPE(BFloat16, float)
|
| 131 |
+
CUDA_ACC_TYPE(Half, float)
|
| 132 |
+
CUDA_ACC_TYPE(Float8_e5m2, float)
|
| 133 |
+
CUDA_ACC_TYPE(Float8_e4m3fn, float)
|
| 134 |
+
CUDA_ACC_TYPE(Float8_e5m2fnuz, float)
|
| 135 |
+
CUDA_ACC_TYPE(Float8_e4m3fnuz, float)
|
| 136 |
+
CUDA_ACC_TYPE(float, float)
|
| 137 |
+
CUDA_ACC_TYPE(double, double)
|
| 138 |
+
CUDA_ACC_TYPE(int8_t, int64_t)
|
| 139 |
+
CUDA_ACC_TYPE(uint8_t, int64_t)
|
| 140 |
+
CUDA_ACC_TYPE(char, int64_t)
|
| 141 |
+
CUDA_ACC_TYPE(int16_t, int64_t)
|
| 142 |
+
CUDA_ACC_TYPE(int32_t, int64_t)
|
| 143 |
+
CUDA_ACC_TYPE(int64_t, int64_t)
|
| 144 |
+
CUDA_ACC_TYPE(bool, bool)
|
| 145 |
+
CUDA_ACC_TYPE(c10::complex<Half>, c10::complex<float>)
|
| 146 |
+
CUDA_ACC_TYPE(c10::complex<float>, c10::complex<float>)
|
| 147 |
+
CUDA_ACC_TYPE(c10::complex<double>, c10::complex<double>)
|
| 148 |
+
|
| 149 |
+
CPU_ACC_TYPE(BFloat16, float)
|
| 150 |
+
CPU_ACC_TYPE(Half, float)
|
| 151 |
+
CPU_ACC_TYPE(Float8_e5m2, float)
|
| 152 |
+
CPU_ACC_TYPE(Float8_e4m3fn, float)
|
| 153 |
+
CPU_ACC_TYPE(Float8_e5m2fnuz, float)
|
| 154 |
+
CPU_ACC_TYPE(Float8_e4m3fnuz, float)
|
| 155 |
+
CPU_ACC_TYPE(float, double)
|
| 156 |
+
CPU_ACC_TYPE(double, double)
|
| 157 |
+
CPU_ACC_TYPE(int8_t, int64_t)
|
| 158 |
+
CPU_ACC_TYPE(uint8_t, int64_t)
|
| 159 |
+
CPU_ACC_TYPE(char, int64_t)
|
| 160 |
+
CPU_ACC_TYPE(int16_t, int64_t)
|
| 161 |
+
CPU_ACC_TYPE(int32_t, int64_t)
|
| 162 |
+
CPU_ACC_TYPE(int64_t, int64_t)
|
| 163 |
+
CPU_ACC_TYPE(bool, bool)
|
| 164 |
+
CPU_ACC_TYPE(c10::complex<Half>, c10::complex<float>)
|
| 165 |
+
CPU_ACC_TYPE(c10::complex<float>, c10::complex<double>)
|
| 166 |
+
CPU_ACC_TYPE(c10::complex<double>, c10::complex<double>)
|
| 167 |
+
|
| 168 |
+
TORCH_API c10::ScalarType toAccumulateType(
|
| 169 |
+
c10::ScalarType type,
|
| 170 |
+
c10::DeviceType device);
|
| 171 |
+
TORCH_API c10::ScalarType toAccumulateType(c10::ScalarType type, bool is_cuda);
|
| 172 |
+
|
| 173 |
+
} // namespace at
|
lib/python3.12/site-packages/torch/include/ATen/CPUFunctions.h
ADDED
|
@@ -0,0 +1,29 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#include <ATen/core/TensorBody.h>
|
| 2 |
+
|
| 3 |
+
// TODO Undo all logic introduced for Note [Avoiding Include Cycles In Static Dispatch]
|
| 4 |
+
// Code introduced to avoid cyclic dependency in static dispatch is no longer
|
| 5 |
+
// needed as static dispatch logic is moved from TensorBody.h, which caused cycles in the first place,
|
| 6 |
+
// to Operators.cpp for supporting multiple backends with multiple kernels.
|
| 7 |
+
//
|
| 8 |
+
// Note [Avoiding Include Cycles In Static Dispatch]
|
| 9 |
+
// In order to avoid #include cycles in the static dispatch build, we've carefully split out
|
| 10 |
+
// the static function definition files into {DispatchKey}Functions.h and {DispatchKey}Functions_inl.h.
|
| 11 |
+
//
|
| 12 |
+
// Without this split, the include cycle looks like TensorBody.h -> CPUFunctions.h -> TensorBody.h.
|
| 13 |
+
// - TensorBody.h #includes CPUFunctions.h in the static dispatch build, because the tensor methods
|
| 14 |
+
// all need to call into the fastpath C++ API defined in CPUFunctions.h. The methods are also all
|
| 15 |
+
// directly inlined into TensorBody.h.
|
| 16 |
+
// - CPUFunctions.h #includes TensorBody.h because it contains function declarations for the entire C++ API,
|
| 17 |
+
// which include functions that have defaultable std::optional<Tensor> arguments.
|
| 18 |
+
// That requires knowing the full Tensor class definition.
|
| 19 |
+
//
|
| 20 |
+
// We break the cycle by doing the following:
|
| 21 |
+
// - Split out CPUFunction.h into two files: CPUFunctions.h and CPUFunctions_inl.h
|
| 22 |
+
// - CPUFunction.h is a dummy file that just includes the Tensor class and includes CPUFunctions_inl.,
|
| 23 |
+
// - CPUFunctions_inl.h includes everything else
|
| 24 |
+
// - (only in the static dispatch build) TensorBody.h makes sure to finish defining the Tensor class,
|
| 25 |
+
// and then it includes CPUFunctions_inl.h.
|
| 26 |
+
// - All other files that want the cpu fastpath functions can include CPUFunctions.h directly.
|
| 27 |
+
// - This also means that static dispatch build, CPUFunctions.h only needs to
|
| 28 |
+
// #include TensorBody.h, and it will automatically bring in CPUFunctions_inl.h.
|
| 29 |
+
#include <ATen/CPUFunctions_inl.h>
|
lib/python3.12/site-packages/torch/include/ATen/CUDAFunctions_inl.h
ADDED
|
@@ -0,0 +1,625 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
// @generated by torchgen/gen.py from DispatchKeyFunctions_inl.h
|
| 3 |
+
|
| 4 |
+
// NB: The implementing C++ file is RegisterDispatchKey.cpp
|
| 5 |
+
|
| 6 |
+
// The only #includes we need are for custom classes that have defaults in the C++ API
|
| 7 |
+
#include <c10/core/MemoryFormat.h>
|
| 8 |
+
#include <c10/core/Scalar.h>
|
| 9 |
+
#include <ATen/core/Reduction.h>
|
| 10 |
+
|
| 11 |
+
#if defined(AT_PER_OPERATOR_HEADERS) && defined(TORCH_ASSERT_ONLY_METHOD_OPERATORS)
|
| 12 |
+
#error This change adds a dependency on all pytorch operators, meaning the \
|
| 13 |
+
file will need to be re-compiled every time an operator is changed or added. \
|
| 14 |
+
Consider including a specific operator from \
|
| 15 |
+
<ATen/ops/{my_operator}_cuda_dispatch.h>. \
|
| 16 |
+
See NOTE [TORCH_ASSERT_ONLY_METHOD_OPERATORS].
|
| 17 |
+
#endif
|
| 18 |
+
|
| 19 |
+
#include <ATen/ops/_adaptive_avg_pool2d_cuda_dispatch.h>
|
| 20 |
+
#include <ATen/ops/_adaptive_avg_pool2d_backward_cuda_dispatch.h>
|
| 21 |
+
#include <ATen/ops/_adaptive_avg_pool3d_cuda_dispatch.h>
|
| 22 |
+
#include <ATen/ops/_adaptive_avg_pool3d_backward_cuda_dispatch.h>
|
| 23 |
+
#include <ATen/ops/_addmm_activation_cuda_dispatch.h>
|
| 24 |
+
#include <ATen/ops/_aminmax_cuda_dispatch.h>
|
| 25 |
+
#include <ATen/ops/_amp_foreach_non_finite_check_and_unscale_cuda_dispatch.h>
|
| 26 |
+
#include <ATen/ops/_amp_update_scale_cuda_dispatch.h>
|
| 27 |
+
#include <ATen/ops/_assert_async_cuda_dispatch.h>
|
| 28 |
+
#include <ATen/ops/_batch_norm_with_update_cuda_dispatch.h>
|
| 29 |
+
#include <ATen/ops/_cdist_backward_cuda_dispatch.h>
|
| 30 |
+
#include <ATen/ops/_cdist_forward_cuda_dispatch.h>
|
| 31 |
+
#include <ATen/ops/_cholesky_solve_helper_cuda_dispatch.h>
|
| 32 |
+
#include <ATen/ops/_chunk_cat_cuda_dispatch.h>
|
| 33 |
+
#include <ATen/ops/_compute_linear_combination_cuda_dispatch.h>
|
| 34 |
+
#include <ATen/ops/_conv_depthwise2d_cuda_dispatch.h>
|
| 35 |
+
#include <ATen/ops/_convert_indices_from_coo_to_csr_cuda_dispatch.h>
|
| 36 |
+
#include <ATen/ops/_convert_indices_from_csr_to_coo_cuda_dispatch.h>
|
| 37 |
+
#include <ATen/ops/_convert_weight_to_int4pack_cuda_dispatch.h>
|
| 38 |
+
#include <ATen/ops/_cslt_compress_cuda_dispatch.h>
|
| 39 |
+
#include <ATen/ops/_cslt_sparse_mm_cuda_dispatch.h>
|
| 40 |
+
#include <ATen/ops/_cslt_sparse_mm_search_cuda_dispatch.h>
|
| 41 |
+
#include <ATen/ops/_ctc_loss_cuda_dispatch.h>
|
| 42 |
+
#include <ATen/ops/_ctc_loss_backward_cuda_dispatch.h>
|
| 43 |
+
#include <ATen/ops/_cudnn_ctc_loss_cuda_dispatch.h>
|
| 44 |
+
#include <ATen/ops/_cudnn_init_dropout_state_cuda_dispatch.h>
|
| 45 |
+
#include <ATen/ops/_cudnn_rnn_cuda_dispatch.h>
|
| 46 |
+
#include <ATen/ops/_cudnn_rnn_backward_cuda_dispatch.h>
|
| 47 |
+
#include <ATen/ops/_cudnn_rnn_flatten_weight_cuda_dispatch.h>
|
| 48 |
+
#include <ATen/ops/_cummax_helper_cuda_dispatch.h>
|
| 49 |
+
#include <ATen/ops/_cummin_helper_cuda_dispatch.h>
|
| 50 |
+
#include <ATen/ops/_dirichlet_grad_cuda_dispatch.h>
|
| 51 |
+
#include <ATen/ops/_efficient_attention_backward_cuda_dispatch.h>
|
| 52 |
+
#include <ATen/ops/_efficient_attention_forward_cuda_dispatch.h>
|
| 53 |
+
#include <ATen/ops/_efficientzerotensor_cuda_dispatch.h>
|
| 54 |
+
#include <ATen/ops/_embedding_bag_cuda_dispatch.h>
|
| 55 |
+
#include <ATen/ops/_embedding_bag_backward_cuda_dispatch.h>
|
| 56 |
+
#include <ATen/ops/_embedding_bag_dense_backward_cuda_dispatch.h>
|
| 57 |
+
#include <ATen/ops/_embedding_bag_forward_only_cuda_dispatch.h>
|
| 58 |
+
#include <ATen/ops/_embedding_bag_per_sample_weights_backward_cuda_dispatch.h>
|
| 59 |
+
#include <ATen/ops/_fake_quantize_learnable_per_channel_affine_cuda_dispatch.h>
|
| 60 |
+
#include <ATen/ops/_fake_quantize_learnable_per_channel_affine_backward_cuda_dispatch.h>
|
| 61 |
+
#include <ATen/ops/_fake_quantize_learnable_per_tensor_affine_cuda_dispatch.h>
|
| 62 |
+
#include <ATen/ops/_fake_quantize_learnable_per_tensor_affine_backward_cuda_dispatch.h>
|
| 63 |
+
#include <ATen/ops/_fake_quantize_per_tensor_affine_cachemask_tensor_qparams_cuda_dispatch.h>
|
| 64 |
+
#include <ATen/ops/_fft_c2c_cuda_dispatch.h>
|
| 65 |
+
#include <ATen/ops/_fft_c2r_cuda_dispatch.h>
|
| 66 |
+
#include <ATen/ops/_fft_r2c_cuda_dispatch.h>
|
| 67 |
+
#include <ATen/ops/_fill_mem_eff_dropout_mask_cuda_dispatch.h>
|
| 68 |
+
#include <ATen/ops/_flash_attention_backward_cuda_dispatch.h>
|
| 69 |
+
#include <ATen/ops/_flash_attention_forward_cuda_dispatch.h>
|
| 70 |
+
#include <ATen/ops/_foreach_abs_cuda_dispatch.h>
|
| 71 |
+
#include <ATen/ops/_foreach_acos_cuda_dispatch.h>
|
| 72 |
+
#include <ATen/ops/_foreach_add_cuda_dispatch.h>
|
| 73 |
+
#include <ATen/ops/_foreach_addcdiv_cuda_dispatch.h>
|
| 74 |
+
#include <ATen/ops/_foreach_addcmul_cuda_dispatch.h>
|
| 75 |
+
#include <ATen/ops/_foreach_asin_cuda_dispatch.h>
|
| 76 |
+
#include <ATen/ops/_foreach_atan_cuda_dispatch.h>
|
| 77 |
+
#include <ATen/ops/_foreach_ceil_cuda_dispatch.h>
|
| 78 |
+
#include <ATen/ops/_foreach_clamp_max_cuda_dispatch.h>
|
| 79 |
+
#include <ATen/ops/_foreach_clamp_min_cuda_dispatch.h>
|
| 80 |
+
#include <ATen/ops/_foreach_copy_cuda_dispatch.h>
|
| 81 |
+
#include <ATen/ops/_foreach_cos_cuda_dispatch.h>
|
| 82 |
+
#include <ATen/ops/_foreach_cosh_cuda_dispatch.h>
|
| 83 |
+
#include <ATen/ops/_foreach_div_cuda_dispatch.h>
|
| 84 |
+
#include <ATen/ops/_foreach_erf_cuda_dispatch.h>
|
| 85 |
+
#include <ATen/ops/_foreach_erfc_cuda_dispatch.h>
|
| 86 |
+
#include <ATen/ops/_foreach_exp_cuda_dispatch.h>
|
| 87 |
+
#include <ATen/ops/_foreach_expm1_cuda_dispatch.h>
|
| 88 |
+
#include <ATen/ops/_foreach_floor_cuda_dispatch.h>
|
| 89 |
+
#include <ATen/ops/_foreach_frac_cuda_dispatch.h>
|
| 90 |
+
#include <ATen/ops/_foreach_lerp_cuda_dispatch.h>
|
| 91 |
+
#include <ATen/ops/_foreach_lgamma_cuda_dispatch.h>
|
| 92 |
+
#include <ATen/ops/_foreach_log_cuda_dispatch.h>
|
| 93 |
+
#include <ATen/ops/_foreach_log10_cuda_dispatch.h>
|
| 94 |
+
#include <ATen/ops/_foreach_log1p_cuda_dispatch.h>
|
| 95 |
+
#include <ATen/ops/_foreach_log2_cuda_dispatch.h>
|
| 96 |
+
#include <ATen/ops/_foreach_max_cuda_dispatch.h>
|
| 97 |
+
#include <ATen/ops/_foreach_maximum_cuda_dispatch.h>
|
| 98 |
+
#include <ATen/ops/_foreach_minimum_cuda_dispatch.h>
|
| 99 |
+
#include <ATen/ops/_foreach_mul_cuda_dispatch.h>
|
| 100 |
+
#include <ATen/ops/_foreach_neg_cuda_dispatch.h>
|
| 101 |
+
#include <ATen/ops/_foreach_norm_cuda_dispatch.h>
|
| 102 |
+
#include <ATen/ops/_foreach_pow_cuda_dispatch.h>
|
| 103 |
+
#include <ATen/ops/_foreach_reciprocal_cuda_dispatch.h>
|
| 104 |
+
#include <ATen/ops/_foreach_round_cuda_dispatch.h>
|
| 105 |
+
#include <ATen/ops/_foreach_rsqrt_cuda_dispatch.h>
|
| 106 |
+
#include <ATen/ops/_foreach_sigmoid_cuda_dispatch.h>
|
| 107 |
+
#include <ATen/ops/_foreach_sign_cuda_dispatch.h>
|
| 108 |
+
#include <ATen/ops/_foreach_sin_cuda_dispatch.h>
|
| 109 |
+
#include <ATen/ops/_foreach_sinh_cuda_dispatch.h>
|
| 110 |
+
#include <ATen/ops/_foreach_sqrt_cuda_dispatch.h>
|
| 111 |
+
#include <ATen/ops/_foreach_sub_cuda_dispatch.h>
|
| 112 |
+
#include <ATen/ops/_foreach_tan_cuda_dispatch.h>
|
| 113 |
+
#include <ATen/ops/_foreach_tanh_cuda_dispatch.h>
|
| 114 |
+
#include <ATen/ops/_foreach_trunc_cuda_dispatch.h>
|
| 115 |
+
#include <ATen/ops/_foreach_zero_cuda_dispatch.h>
|
| 116 |
+
#include <ATen/ops/_fused_adam_cuda_dispatch.h>
|
| 117 |
+
#include <ATen/ops/_fused_adamw_cuda_dispatch.h>
|
| 118 |
+
#include <ATen/ops/_fused_dropout_cuda_dispatch.h>
|
| 119 |
+
#include <ATen/ops/_fused_moving_avg_obs_fq_helper_cuda_dispatch.h>
|
| 120 |
+
#include <ATen/ops/_fused_sdp_choice_cuda_dispatch.h>
|
| 121 |
+
#include <ATen/ops/_fused_sgd_cuda_dispatch.h>
|
| 122 |
+
#include <ATen/ops/_index_put_impl_cuda_dispatch.h>
|
| 123 |
+
#include <ATen/ops/_int_mm_cuda_dispatch.h>
|
| 124 |
+
#include <ATen/ops/_jagged_to_padded_dense_forward_cuda_dispatch.h>
|
| 125 |
+
#include <ATen/ops/_linalg_det_cuda_dispatch.h>
|
| 126 |
+
#include <ATen/ops/_linalg_eigh_cuda_dispatch.h>
|
| 127 |
+
#include <ATen/ops/_linalg_eigvals_cuda_dispatch.h>
|
| 128 |
+
#include <ATen/ops/_linalg_slogdet_cuda_dispatch.h>
|
| 129 |
+
#include <ATen/ops/_linalg_solve_ex_cuda_dispatch.h>
|
| 130 |
+
#include <ATen/ops/_linalg_svd_cuda_dispatch.h>
|
| 131 |
+
#include <ATen/ops/_local_scalar_dense_cuda_dispatch.h>
|
| 132 |
+
#include <ATen/ops/_log_softmax_cuda_dispatch.h>
|
| 133 |
+
#include <ATen/ops/_log_softmax_backward_data_cuda_dispatch.h>
|
| 134 |
+
#include <ATen/ops/_logcumsumexp_cuda_dispatch.h>
|
| 135 |
+
#include <ATen/ops/_make_per_channel_quantized_tensor_cuda_dispatch.h>
|
| 136 |
+
#include <ATen/ops/_make_per_tensor_quantized_tensor_cuda_dispatch.h>
|
| 137 |
+
#include <ATen/ops/_masked_scale_cuda_dispatch.h>
|
| 138 |
+
#include <ATen/ops/_masked_softmax_cuda_dispatch.h>
|
| 139 |
+
#include <ATen/ops/_masked_softmax_backward_cuda_dispatch.h>
|
| 140 |
+
#include <ATen/ops/_mixed_dtypes_linear_cuda_dispatch.h>
|
| 141 |
+
#include <ATen/ops/_native_batch_norm_legit_cuda_dispatch.h>
|
| 142 |
+
#include <ATen/ops/_native_multi_head_attention_cuda_dispatch.h>
|
| 143 |
+
#include <ATen/ops/_nested_compute_contiguous_strides_offsets_cuda_dispatch.h>
|
| 144 |
+
#include <ATen/ops/_nested_from_padded_cuda_dispatch.h>
|
| 145 |
+
#include <ATen/ops/_nested_tensor_from_mask_cuda_dispatch.h>
|
| 146 |
+
#include <ATen/ops/_nested_tensor_from_mask_left_aligned_cuda_dispatch.h>
|
| 147 |
+
#include <ATen/ops/_nested_view_from_buffer_cuda_dispatch.h>
|
| 148 |
+
#include <ATen/ops/_padded_dense_to_jagged_forward_cuda_dispatch.h>
|
| 149 |
+
#include <ATen/ops/_pdist_backward_cuda_dispatch.h>
|
| 150 |
+
#include <ATen/ops/_pdist_forward_cuda_dispatch.h>
|
| 151 |
+
#include <ATen/ops/_prelu_kernel_cuda_dispatch.h>
|
| 152 |
+
#include <ATen/ops/_prelu_kernel_backward_cuda_dispatch.h>
|
| 153 |
+
#include <ATen/ops/_reshape_alias_cuda_dispatch.h>
|
| 154 |
+
#include <ATen/ops/_sample_dirichlet_cuda_dispatch.h>
|
| 155 |
+
#include <ATen/ops/_scaled_dot_product_cudnn_attention_cuda_dispatch.h>
|
| 156 |
+
#include <ATen/ops/_scaled_dot_product_cudnn_attention_backward_cuda_dispatch.h>
|
| 157 |
+
#include <ATen/ops/_scaled_dot_product_efficient_attention_cuda_dispatch.h>
|
| 158 |
+
#include <ATen/ops/_scaled_dot_product_efficient_attention_backward_cuda_dispatch.h>
|
| 159 |
+
#include <ATen/ops/_scaled_dot_product_flash_attention_cuda_dispatch.h>
|
| 160 |
+
#include <ATen/ops/_scaled_dot_product_flash_attention_backward_cuda_dispatch.h>
|
| 161 |
+
#include <ATen/ops/_scaled_mm_cuda_dispatch.h>
|
| 162 |
+
#include <ATen/ops/_segment_reduce_backward_cuda_dispatch.h>
|
| 163 |
+
#include <ATen/ops/_slow_conv2d_backward_cuda_dispatch.h>
|
| 164 |
+
#include <ATen/ops/_slow_conv2d_forward_cuda_dispatch.h>
|
| 165 |
+
#include <ATen/ops/_softmax_cuda_dispatch.h>
|
| 166 |
+
#include <ATen/ops/_softmax_backward_data_cuda_dispatch.h>
|
| 167 |
+
#include <ATen/ops/_sparse_semi_structured_addmm_cuda_dispatch.h>
|
| 168 |
+
#include <ATen/ops/_sparse_semi_structured_apply_cuda_dispatch.h>
|
| 169 |
+
#include <ATen/ops/_sparse_semi_structured_apply_dense_cuda_dispatch.h>
|
| 170 |
+
#include <ATen/ops/_sparse_semi_structured_linear_cuda_dispatch.h>
|
| 171 |
+
#include <ATen/ops/_sparse_semi_structured_mm_cuda_dispatch.h>
|
| 172 |
+
#include <ATen/ops/_sparse_semi_structured_tile_cuda_dispatch.h>
|
| 173 |
+
#include <ATen/ops/_standard_gamma_cuda_dispatch.h>
|
| 174 |
+
#include <ATen/ops/_standard_gamma_grad_cuda_dispatch.h>
|
| 175 |
+
#include <ATen/ops/_thnn_fused_gru_cell_cuda_dispatch.h>
|
| 176 |
+
#include <ATen/ops/_thnn_fused_gru_cell_backward_cuda_dispatch.h>
|
| 177 |
+
#include <ATen/ops/_thnn_fused_lstm_cell_cuda_dispatch.h>
|
| 178 |
+
#include <ATen/ops/_thnn_fused_lstm_cell_backward_impl_cuda_dispatch.h>
|
| 179 |
+
#include <ATen/ops/_to_sparse_cuda_dispatch.h>
|
| 180 |
+
#include <ATen/ops/_to_sparse_bsc_cuda_dispatch.h>
|
| 181 |
+
#include <ATen/ops/_to_sparse_bsr_cuda_dispatch.h>
|
| 182 |
+
#include <ATen/ops/_to_sparse_csc_cuda_dispatch.h>
|
| 183 |
+
#include <ATen/ops/_to_sparse_csr_cuda_dispatch.h>
|
| 184 |
+
#include <ATen/ops/_to_sparse_semi_structured_cuda_dispatch.h>
|
| 185 |
+
#include <ATen/ops/_transform_bias_rescale_qkv_cuda_dispatch.h>
|
| 186 |
+
#include <ATen/ops/_transformer_encoder_layer_fwd_cuda_dispatch.h>
|
| 187 |
+
#include <ATen/ops/_triton_multi_head_attention_cuda_dispatch.h>
|
| 188 |
+
#include <ATen/ops/_triton_scaled_dot_attention_cuda_dispatch.h>
|
| 189 |
+
#include <ATen/ops/_unique_cuda_dispatch.h>
|
| 190 |
+
#include <ATen/ops/_unique2_cuda_dispatch.h>
|
| 191 |
+
#include <ATen/ops/_upsample_bicubic2d_aa_cuda_dispatch.h>
|
| 192 |
+
#include <ATen/ops/_upsample_bicubic2d_aa_backward_cuda_dispatch.h>
|
| 193 |
+
#include <ATen/ops/_upsample_bilinear2d_aa_cuda_dispatch.h>
|
| 194 |
+
#include <ATen/ops/_upsample_bilinear2d_aa_backward_cuda_dispatch.h>
|
| 195 |
+
#include <ATen/ops/_upsample_nearest_exact1d_cuda_dispatch.h>
|
| 196 |
+
#include <ATen/ops/_upsample_nearest_exact1d_backward_cuda_dispatch.h>
|
| 197 |
+
#include <ATen/ops/_upsample_nearest_exact2d_cuda_dispatch.h>
|
| 198 |
+
#include <ATen/ops/_upsample_nearest_exact2d_backward_cuda_dispatch.h>
|
| 199 |
+
#include <ATen/ops/_upsample_nearest_exact3d_cuda_dispatch.h>
|
| 200 |
+
#include <ATen/ops/_upsample_nearest_exact3d_backward_cuda_dispatch.h>
|
| 201 |
+
#include <ATen/ops/_use_cudnn_ctc_loss_cuda_dispatch.h>
|
| 202 |
+
#include <ATen/ops/_validate_compressed_sparse_indices_cuda_dispatch.h>
|
| 203 |
+
#include <ATen/ops/_weight_int4pack_mm_cuda_dispatch.h>
|
| 204 |
+
#include <ATen/ops/_weight_norm_interface_cuda_dispatch.h>
|
| 205 |
+
#include <ATen/ops/_weight_norm_interface_backward_cuda_dispatch.h>
|
| 206 |
+
#include <ATen/ops/abs_cuda_dispatch.h>
|
| 207 |
+
#include <ATen/ops/acos_cuda_dispatch.h>
|
| 208 |
+
#include <ATen/ops/acosh_cuda_dispatch.h>
|
| 209 |
+
#include <ATen/ops/adaptive_avg_pool2d_cuda_dispatch.h>
|
| 210 |
+
#include <ATen/ops/adaptive_avg_pool3d_cuda_dispatch.h>
|
| 211 |
+
#include <ATen/ops/adaptive_avg_pool3d_backward_cuda_dispatch.h>
|
| 212 |
+
#include <ATen/ops/adaptive_max_pool2d_cuda_dispatch.h>
|
| 213 |
+
#include <ATen/ops/adaptive_max_pool2d_backward_cuda_dispatch.h>
|
| 214 |
+
#include <ATen/ops/adaptive_max_pool3d_cuda_dispatch.h>
|
| 215 |
+
#include <ATen/ops/adaptive_max_pool3d_backward_cuda_dispatch.h>
|
| 216 |
+
#include <ATen/ops/add_cuda_dispatch.h>
|
| 217 |
+
#include <ATen/ops/addbmm_cuda_dispatch.h>
|
| 218 |
+
#include <ATen/ops/addcdiv_cuda_dispatch.h>
|
| 219 |
+
#include <ATen/ops/addcmul_cuda_dispatch.h>
|
| 220 |
+
#include <ATen/ops/addmm_cuda_dispatch.h>
|
| 221 |
+
#include <ATen/ops/addmv_cuda_dispatch.h>
|
| 222 |
+
#include <ATen/ops/addr_cuda_dispatch.h>
|
| 223 |
+
#include <ATen/ops/all_cuda_dispatch.h>
|
| 224 |
+
#include <ATen/ops/amax_cuda_dispatch.h>
|
| 225 |
+
#include <ATen/ops/amin_cuda_dispatch.h>
|
| 226 |
+
#include <ATen/ops/aminmax_cuda_dispatch.h>
|
| 227 |
+
#include <ATen/ops/angle_cuda_dispatch.h>
|
| 228 |
+
#include <ATen/ops/any_cuda_dispatch.h>
|
| 229 |
+
#include <ATen/ops/arange_cuda_dispatch.h>
|
| 230 |
+
#include <ATen/ops/argmax_cuda_dispatch.h>
|
| 231 |
+
#include <ATen/ops/argmin_cuda_dispatch.h>
|
| 232 |
+
#include <ATen/ops/as_strided_cuda_dispatch.h>
|
| 233 |
+
#include <ATen/ops/asin_cuda_dispatch.h>
|
| 234 |
+
#include <ATen/ops/asinh_cuda_dispatch.h>
|
| 235 |
+
#include <ATen/ops/atan_cuda_dispatch.h>
|
| 236 |
+
#include <ATen/ops/atan2_cuda_dispatch.h>
|
| 237 |
+
#include <ATen/ops/atanh_cuda_dispatch.h>
|
| 238 |
+
#include <ATen/ops/avg_pool2d_cuda_dispatch.h>
|
| 239 |
+
#include <ATen/ops/avg_pool2d_backward_cuda_dispatch.h>
|
| 240 |
+
#include <ATen/ops/avg_pool3d_cuda_dispatch.h>
|
| 241 |
+
#include <ATen/ops/avg_pool3d_backward_cuda_dispatch.h>
|
| 242 |
+
#include <ATen/ops/baddbmm_cuda_dispatch.h>
|
| 243 |
+
#include <ATen/ops/batch_norm_backward_cuda_dispatch.h>
|
| 244 |
+
#include <ATen/ops/batch_norm_backward_elemt_cuda_dispatch.h>
|
| 245 |
+
#include <ATen/ops/batch_norm_backward_reduce_cuda_dispatch.h>
|
| 246 |
+
#include <ATen/ops/batch_norm_elemt_cuda_dispatch.h>
|
| 247 |
+
#include <ATen/ops/batch_norm_gather_stats_cuda_dispatch.h>
|
| 248 |
+
#include <ATen/ops/batch_norm_gather_stats_with_counts_cuda_dispatch.h>
|
| 249 |
+
#include <ATen/ops/batch_norm_stats_cuda_dispatch.h>
|
| 250 |
+
#include <ATen/ops/batch_norm_update_stats_cuda_dispatch.h>
|
| 251 |
+
#include <ATen/ops/bernoulli_cuda_dispatch.h>
|
| 252 |
+
#include <ATen/ops/binary_cross_entropy_cuda_dispatch.h>
|
| 253 |
+
#include <ATen/ops/binary_cross_entropy_backward_cuda_dispatch.h>
|
| 254 |
+
#include <ATen/ops/bincount_cuda_dispatch.h>
|
| 255 |
+
#include <ATen/ops/binomial_cuda_dispatch.h>
|
| 256 |
+
#include <ATen/ops/bitwise_and_cuda_dispatch.h>
|
| 257 |
+
#include <ATen/ops/bitwise_left_shift_cuda_dispatch.h>
|
| 258 |
+
#include <ATen/ops/bitwise_not_cuda_dispatch.h>
|
| 259 |
+
#include <ATen/ops/bitwise_or_cuda_dispatch.h>
|
| 260 |
+
#include <ATen/ops/bitwise_right_shift_cuda_dispatch.h>
|
| 261 |
+
#include <ATen/ops/bitwise_xor_cuda_dispatch.h>
|
| 262 |
+
#include <ATen/ops/bmm_cuda_dispatch.h>
|
| 263 |
+
#include <ATen/ops/bucketize_cuda_dispatch.h>
|
| 264 |
+
#include <ATen/ops/cat_cuda_dispatch.h>
|
| 265 |
+
#include <ATen/ops/cauchy_cuda_dispatch.h>
|
| 266 |
+
#include <ATen/ops/ceil_cuda_dispatch.h>
|
| 267 |
+
#include <ATen/ops/channel_shuffle_cuda_dispatch.h>
|
| 268 |
+
#include <ATen/ops/cholesky_cuda_dispatch.h>
|
| 269 |
+
#include <ATen/ops/cholesky_inverse_cuda_dispatch.h>
|
| 270 |
+
#include <ATen/ops/clamp_cuda_dispatch.h>
|
| 271 |
+
#include <ATen/ops/clamp_max_cuda_dispatch.h>
|
| 272 |
+
#include <ATen/ops/clamp_min_cuda_dispatch.h>
|
| 273 |
+
#include <ATen/ops/col2im_cuda_dispatch.h>
|
| 274 |
+
#include <ATen/ops/complex_cuda_dispatch.h>
|
| 275 |
+
#include <ATen/ops/conj_physical_cuda_dispatch.h>
|
| 276 |
+
#include <ATen/ops/conv_depthwise3d_cuda_dispatch.h>
|
| 277 |
+
#include <ATen/ops/convolution_backward_cuda_dispatch.h>
|
| 278 |
+
#include <ATen/ops/copysign_cuda_dispatch.h>
|
| 279 |
+
#include <ATen/ops/cos_cuda_dispatch.h>
|
| 280 |
+
#include <ATen/ops/cosh_cuda_dispatch.h>
|
| 281 |
+
#include <ATen/ops/count_nonzero_cuda_dispatch.h>
|
| 282 |
+
#include <ATen/ops/cudnn_affine_grid_generator_cuda_dispatch.h>
|
| 283 |
+
#include <ATen/ops/cudnn_affine_grid_generator_backward_cuda_dispatch.h>
|
| 284 |
+
#include <ATen/ops/cudnn_batch_norm_cuda_dispatch.h>
|
| 285 |
+
#include <ATen/ops/cudnn_batch_norm_backward_cuda_dispatch.h>
|
| 286 |
+
#include <ATen/ops/cudnn_convolution_cuda_dispatch.h>
|
| 287 |
+
#include <ATen/ops/cudnn_convolution_add_relu_cuda_dispatch.h>
|
| 288 |
+
#include <ATen/ops/cudnn_convolution_relu_cuda_dispatch.h>
|
| 289 |
+
#include <ATen/ops/cudnn_convolution_transpose_cuda_dispatch.h>
|
| 290 |
+
#include <ATen/ops/cudnn_grid_sampler_cuda_dispatch.h>
|
| 291 |
+
#include <ATen/ops/cudnn_grid_sampler_backward_cuda_dispatch.h>
|
| 292 |
+
#include <ATen/ops/cumprod_cuda_dispatch.h>
|
| 293 |
+
#include <ATen/ops/cumsum_cuda_dispatch.h>
|
| 294 |
+
#include <ATen/ops/dequantize_cuda_dispatch.h>
|
| 295 |
+
#include <ATen/ops/digamma_cuda_dispatch.h>
|
| 296 |
+
#include <ATen/ops/div_cuda_dispatch.h>
|
| 297 |
+
#include <ATen/ops/dot_cuda_dispatch.h>
|
| 298 |
+
#include <ATen/ops/elu_cuda_dispatch.h>
|
| 299 |
+
#include <ATen/ops/elu_backward_cuda_dispatch.h>
|
| 300 |
+
#include <ATen/ops/embedding_dense_backward_cuda_dispatch.h>
|
| 301 |
+
#include <ATen/ops/embedding_renorm_cuda_dispatch.h>
|
| 302 |
+
#include <ATen/ops/empty_cuda_dispatch.h>
|
| 303 |
+
#include <ATen/ops/empty_strided_cuda_dispatch.h>
|
| 304 |
+
#include <ATen/ops/eq_cuda_dispatch.h>
|
| 305 |
+
#include <ATen/ops/equal_cuda_dispatch.h>
|
| 306 |
+
#include <ATen/ops/erf_cuda_dispatch.h>
|
| 307 |
+
#include <ATen/ops/erfc_cuda_dispatch.h>
|
| 308 |
+
#include <ATen/ops/erfinv_cuda_dispatch.h>
|
| 309 |
+
#include <ATen/ops/exp_cuda_dispatch.h>
|
| 310 |
+
#include <ATen/ops/exp2_cuda_dispatch.h>
|
| 311 |
+
#include <ATen/ops/expm1_cuda_dispatch.h>
|
| 312 |
+
#include <ATen/ops/exponential_cuda_dispatch.h>
|
| 313 |
+
#include <ATen/ops/eye_cuda_dispatch.h>
|
| 314 |
+
#include <ATen/ops/fake_quantize_per_channel_affine_cachemask_cuda_dispatch.h>
|
| 315 |
+
#include <ATen/ops/fake_quantize_per_tensor_affine_cachemask_cuda_dispatch.h>
|
| 316 |
+
#include <ATen/ops/fill_cuda_dispatch.h>
|
| 317 |
+
#include <ATen/ops/flip_cuda_dispatch.h>
|
| 318 |
+
#include <ATen/ops/floor_cuda_dispatch.h>
|
| 319 |
+
#include <ATen/ops/floor_divide_cuda_dispatch.h>
|
| 320 |
+
#include <ATen/ops/fmax_cuda_dispatch.h>
|
| 321 |
+
#include <ATen/ops/fmin_cuda_dispatch.h>
|
| 322 |
+
#include <ATen/ops/fmod_cuda_dispatch.h>
|
| 323 |
+
#include <ATen/ops/frac_cuda_dispatch.h>
|
| 324 |
+
#include <ATen/ops/fractional_max_pool2d_cuda_dispatch.h>
|
| 325 |
+
#include <ATen/ops/fractional_max_pool2d_backward_cuda_dispatch.h>
|
| 326 |
+
#include <ATen/ops/fractional_max_pool3d_cuda_dispatch.h>
|
| 327 |
+
#include <ATen/ops/fractional_max_pool3d_backward_cuda_dispatch.h>
|
| 328 |
+
#include <ATen/ops/frexp_cuda_dispatch.h>
|
| 329 |
+
#include <ATen/ops/gather_cuda_dispatch.h>
|
| 330 |
+
#include <ATen/ops/gcd_cuda_dispatch.h>
|
| 331 |
+
#include <ATen/ops/ge_cuda_dispatch.h>
|
| 332 |
+
#include <ATen/ops/gelu_cuda_dispatch.h>
|
| 333 |
+
#include <ATen/ops/gelu_backward_cuda_dispatch.h>
|
| 334 |
+
#include <ATen/ops/geometric_cuda_dispatch.h>
|
| 335 |
+
#include <ATen/ops/geqrf_cuda_dispatch.h>
|
| 336 |
+
#include <ATen/ops/glu_cuda_dispatch.h>
|
| 337 |
+
#include <ATen/ops/glu_backward_cuda_dispatch.h>
|
| 338 |
+
#include <ATen/ops/glu_backward_jvp_cuda_dispatch.h>
|
| 339 |
+
#include <ATen/ops/glu_jvp_cuda_dispatch.h>
|
| 340 |
+
#include <ATen/ops/grid_sampler_2d_cuda_dispatch.h>
|
| 341 |
+
#include <ATen/ops/grid_sampler_2d_backward_cuda_dispatch.h>
|
| 342 |
+
#include <ATen/ops/grid_sampler_3d_cuda_dispatch.h>
|
| 343 |
+
#include <ATen/ops/grid_sampler_3d_backward_cuda_dispatch.h>
|
| 344 |
+
#include <ATen/ops/gt_cuda_dispatch.h>
|
| 345 |
+
#include <ATen/ops/hardshrink_cuda_dispatch.h>
|
| 346 |
+
#include <ATen/ops/hardshrink_backward_cuda_dispatch.h>
|
| 347 |
+
#include <ATen/ops/hardsigmoid_cuda_dispatch.h>
|
| 348 |
+
#include <ATen/ops/hardsigmoid_backward_cuda_dispatch.h>
|
| 349 |
+
#include <ATen/ops/hardswish_cuda_dispatch.h>
|
| 350 |
+
#include <ATen/ops/hardswish_backward_cuda_dispatch.h>
|
| 351 |
+
#include <ATen/ops/hardtanh_cuda_dispatch.h>
|
| 352 |
+
#include <ATen/ops/hardtanh_backward_cuda_dispatch.h>
|
| 353 |
+
#include <ATen/ops/heaviside_cuda_dispatch.h>
|
| 354 |
+
#include <ATen/ops/histc_cuda_dispatch.h>
|
| 355 |
+
#include <ATen/ops/huber_loss_cuda_dispatch.h>
|
| 356 |
+
#include <ATen/ops/huber_loss_backward_cuda_dispatch.h>
|
| 357 |
+
#include <ATen/ops/hypot_cuda_dispatch.h>
|
| 358 |
+
#include <ATen/ops/i0_cuda_dispatch.h>
|
| 359 |
+
#include <ATen/ops/igamma_cuda_dispatch.h>
|
| 360 |
+
#include <ATen/ops/igammac_cuda_dispatch.h>
|
| 361 |
+
#include <ATen/ops/im2col_cuda_dispatch.h>
|
| 362 |
+
#include <ATen/ops/index_cuda_dispatch.h>
|
| 363 |
+
#include <ATen/ops/index_add_cuda_dispatch.h>
|
| 364 |
+
#include <ATen/ops/index_copy_cuda_dispatch.h>
|
| 365 |
+
#include <ATen/ops/index_fill_cuda_dispatch.h>
|
| 366 |
+
#include <ATen/ops/index_reduce_cuda_dispatch.h>
|
| 367 |
+
#include <ATen/ops/index_select_cuda_dispatch.h>
|
| 368 |
+
#include <ATen/ops/is_set_to_cuda_dispatch.h>
|
| 369 |
+
#include <ATen/ops/isin_cuda_dispatch.h>
|
| 370 |
+
#include <ATen/ops/isnan_cuda_dispatch.h>
|
| 371 |
+
#include <ATen/ops/isneginf_cuda_dispatch.h>
|
| 372 |
+
#include <ATen/ops/isposinf_cuda_dispatch.h>
|
| 373 |
+
#include <ATen/ops/kthvalue_cuda_dispatch.h>
|
| 374 |
+
#include <ATen/ops/lcm_cuda_dispatch.h>
|
| 375 |
+
#include <ATen/ops/le_cuda_dispatch.h>
|
| 376 |
+
#include <ATen/ops/leaky_relu_cuda_dispatch.h>
|
| 377 |
+
#include <ATen/ops/leaky_relu_backward_cuda_dispatch.h>
|
| 378 |
+
#include <ATen/ops/lerp_cuda_dispatch.h>
|
| 379 |
+
#include <ATen/ops/lgamma_cuda_dispatch.h>
|
| 380 |
+
#include <ATen/ops/linalg_cholesky_ex_cuda_dispatch.h>
|
| 381 |
+
#include <ATen/ops/linalg_cross_cuda_dispatch.h>
|
| 382 |
+
#include <ATen/ops/linalg_eig_cuda_dispatch.h>
|
| 383 |
+
#include <ATen/ops/linalg_eigvals_cuda_dispatch.h>
|
| 384 |
+
#include <ATen/ops/linalg_householder_product_cuda_dispatch.h>
|
| 385 |
+
#include <ATen/ops/linalg_inv_ex_cuda_dispatch.h>
|
| 386 |
+
#include <ATen/ops/linalg_ldl_factor_ex_cuda_dispatch.h>
|
| 387 |
+
#include <ATen/ops/linalg_ldl_solve_cuda_dispatch.h>
|
| 388 |
+
#include <ATen/ops/linalg_lstsq_cuda_dispatch.h>
|
| 389 |
+
#include <ATen/ops/linalg_lu_cuda_dispatch.h>
|
| 390 |
+
#include <ATen/ops/linalg_lu_factor_ex_cuda_dispatch.h>
|
| 391 |
+
#include <ATen/ops/linalg_lu_solve_cuda_dispatch.h>
|
| 392 |
+
#include <ATen/ops/linalg_matrix_exp_cuda_dispatch.h>
|
| 393 |
+
#include <ATen/ops/linalg_qr_cuda_dispatch.h>
|
| 394 |
+
#include <ATen/ops/linalg_solve_triangular_cuda_dispatch.h>
|
| 395 |
+
#include <ATen/ops/linalg_vector_norm_cuda_dispatch.h>
|
| 396 |
+
#include <ATen/ops/linspace_cuda_dispatch.h>
|
| 397 |
+
#include <ATen/ops/log_cuda_dispatch.h>
|
| 398 |
+
#include <ATen/ops/log10_cuda_dispatch.h>
|
| 399 |
+
#include <ATen/ops/log1p_cuda_dispatch.h>
|
| 400 |
+
#include <ATen/ops/log2_cuda_dispatch.h>
|
| 401 |
+
#include <ATen/ops/log_normal_cuda_dispatch.h>
|
| 402 |
+
#include <ATen/ops/log_sigmoid_backward_cuda_dispatch.h>
|
| 403 |
+
#include <ATen/ops/log_sigmoid_forward_cuda_dispatch.h>
|
| 404 |
+
#include <ATen/ops/logaddexp_cuda_dispatch.h>
|
| 405 |
+
#include <ATen/ops/logaddexp2_cuda_dispatch.h>
|
| 406 |
+
#include <ATen/ops/logical_and_cuda_dispatch.h>
|
| 407 |
+
#include <ATen/ops/logical_not_cuda_dispatch.h>
|
| 408 |
+
#include <ATen/ops/logical_or_cuda_dispatch.h>
|
| 409 |
+
#include <ATen/ops/logical_xor_cuda_dispatch.h>
|
| 410 |
+
#include <ATen/ops/logit_cuda_dispatch.h>
|
| 411 |
+
#include <ATen/ops/logit_backward_cuda_dispatch.h>
|
| 412 |
+
#include <ATen/ops/logspace_cuda_dispatch.h>
|
| 413 |
+
#include <ATen/ops/lshift_cuda_dispatch.h>
|
| 414 |
+
#include <ATen/ops/lt_cuda_dispatch.h>
|
| 415 |
+
#include <ATen/ops/lu_unpack_cuda_dispatch.h>
|
| 416 |
+
#include <ATen/ops/masked_fill_cuda_dispatch.h>
|
| 417 |
+
#include <ATen/ops/masked_scatter_cuda_dispatch.h>
|
| 418 |
+
#include <ATen/ops/masked_select_cuda_dispatch.h>
|
| 419 |
+
#include <ATen/ops/max_cuda_dispatch.h>
|
| 420 |
+
#include <ATen/ops/max_pool2d_with_indices_cuda_dispatch.h>
|
| 421 |
+
#include <ATen/ops/max_pool2d_with_indices_backward_cuda_dispatch.h>
|
| 422 |
+
#include <ATen/ops/max_pool3d_with_indices_cuda_dispatch.h>
|
| 423 |
+
#include <ATen/ops/max_pool3d_with_indices_backward_cuda_dispatch.h>
|
| 424 |
+
#include <ATen/ops/max_unpool2d_cuda_dispatch.h>
|
| 425 |
+
#include <ATen/ops/max_unpool3d_cuda_dispatch.h>
|
| 426 |
+
#include <ATen/ops/maximum_cuda_dispatch.h>
|
| 427 |
+
#include <ATen/ops/mean_cuda_dispatch.h>
|
| 428 |
+
#include <ATen/ops/median_cuda_dispatch.h>
|
| 429 |
+
#include <ATen/ops/min_cuda_dispatch.h>
|
| 430 |
+
#include <ATen/ops/minimum_cuda_dispatch.h>
|
| 431 |
+
#include <ATen/ops/miopen_batch_norm_cuda_dispatch.h>
|
| 432 |
+
#include <ATen/ops/miopen_batch_norm_backward_cuda_dispatch.h>
|
| 433 |
+
#include <ATen/ops/miopen_convolution_cuda_dispatch.h>
|
| 434 |
+
#include <ATen/ops/miopen_convolution_add_relu_cuda_dispatch.h>
|
| 435 |
+
#include <ATen/ops/miopen_convolution_relu_cuda_dispatch.h>
|
| 436 |
+
#include <ATen/ops/miopen_convolution_transpose_cuda_dispatch.h>
|
| 437 |
+
#include <ATen/ops/miopen_depthwise_convolution_cuda_dispatch.h>
|
| 438 |
+
#include <ATen/ops/miopen_rnn_cuda_dispatch.h>
|
| 439 |
+
#include <ATen/ops/miopen_rnn_backward_cuda_dispatch.h>
|
| 440 |
+
#include <ATen/ops/mish_cuda_dispatch.h>
|
| 441 |
+
#include <ATen/ops/mish_backward_cuda_dispatch.h>
|
| 442 |
+
#include <ATen/ops/mm_cuda_dispatch.h>
|
| 443 |
+
#include <ATen/ops/mode_cuda_dispatch.h>
|
| 444 |
+
#include <ATen/ops/mse_loss_cuda_dispatch.h>
|
| 445 |
+
#include <ATen/ops/mse_loss_backward_cuda_dispatch.h>
|
| 446 |
+
#include <ATen/ops/mul_cuda_dispatch.h>
|
| 447 |
+
#include <ATen/ops/multi_margin_loss_cuda_dispatch.h>
|
| 448 |
+
#include <ATen/ops/multi_margin_loss_backward_cuda_dispatch.h>
|
| 449 |
+
#include <ATen/ops/multilabel_margin_loss_backward_cuda_dispatch.h>
|
| 450 |
+
#include <ATen/ops/multilabel_margin_loss_forward_cuda_dispatch.h>
|
| 451 |
+
#include <ATen/ops/multinomial_cuda_dispatch.h>
|
| 452 |
+
#include <ATen/ops/mvlgamma_cuda_dispatch.h>
|
| 453 |
+
#include <ATen/ops/nan_to_num_cuda_dispatch.h>
|
| 454 |
+
#include <ATen/ops/nanmedian_cuda_dispatch.h>
|
| 455 |
+
#include <ATen/ops/nansum_cuda_dispatch.h>
|
| 456 |
+
#include <ATen/ops/native_batch_norm_cuda_dispatch.h>
|
| 457 |
+
#include <ATen/ops/native_batch_norm_backward_cuda_dispatch.h>
|
| 458 |
+
#include <ATen/ops/native_dropout_cuda_dispatch.h>
|
| 459 |
+
#include <ATen/ops/native_dropout_backward_cuda_dispatch.h>
|
| 460 |
+
#include <ATen/ops/native_group_norm_cuda_dispatch.h>
|
| 461 |
+
#include <ATen/ops/native_group_norm_backward_cuda_dispatch.h>
|
| 462 |
+
#include <ATen/ops/native_layer_norm_cuda_dispatch.h>
|
| 463 |
+
#include <ATen/ops/native_layer_norm_backward_cuda_dispatch.h>
|
| 464 |
+
#include <ATen/ops/ne_cuda_dispatch.h>
|
| 465 |
+
#include <ATen/ops/neg_cuda_dispatch.h>
|
| 466 |
+
#include <ATen/ops/nextafter_cuda_dispatch.h>
|
| 467 |
+
#include <ATen/ops/nll_loss2d_backward_cuda_dispatch.h>
|
| 468 |
+
#include <ATen/ops/nll_loss2d_forward_cuda_dispatch.h>
|
| 469 |
+
#include <ATen/ops/nll_loss_backward_cuda_dispatch.h>
|
| 470 |
+
#include <ATen/ops/nll_loss_forward_cuda_dispatch.h>
|
| 471 |
+
#include <ATen/ops/nonzero_cuda_dispatch.h>
|
| 472 |
+
#include <ATen/ops/nonzero_static_cuda_dispatch.h>
|
| 473 |
+
#include <ATen/ops/norm_cuda_dispatch.h>
|
| 474 |
+
#include <ATen/ops/normal_cuda_dispatch.h>
|
| 475 |
+
#include <ATen/ops/ormqr_cuda_dispatch.h>
|
| 476 |
+
#include <ATen/ops/poisson_cuda_dispatch.h>
|
| 477 |
+
#include <ATen/ops/polar_cuda_dispatch.h>
|
| 478 |
+
#include <ATen/ops/polygamma_cuda_dispatch.h>
|
| 479 |
+
#include <ATen/ops/pow_cuda_dispatch.h>
|
| 480 |
+
#include <ATen/ops/prod_cuda_dispatch.h>
|
| 481 |
+
#include <ATen/ops/put_cuda_dispatch.h>
|
| 482 |
+
#include <ATen/ops/quantize_per_channel_cuda_dispatch.h>
|
| 483 |
+
#include <ATen/ops/quantize_per_tensor_cuda_dispatch.h>
|
| 484 |
+
#include <ATen/ops/quantize_per_tensor_dynamic_cuda_dispatch.h>
|
| 485 |
+
#include <ATen/ops/random_cuda_dispatch.h>
|
| 486 |
+
#include <ATen/ops/randperm_cuda_dispatch.h>
|
| 487 |
+
#include <ATen/ops/range_cuda_dispatch.h>
|
| 488 |
+
#include <ATen/ops/reciprocal_cuda_dispatch.h>
|
| 489 |
+
#include <ATen/ops/record_stream_cuda_dispatch.h>
|
| 490 |
+
#include <ATen/ops/reflection_pad1d_cuda_dispatch.h>
|
| 491 |
+
#include <ATen/ops/reflection_pad1d_backward_cuda_dispatch.h>
|
| 492 |
+
#include <ATen/ops/reflection_pad2d_cuda_dispatch.h>
|
| 493 |
+
#include <ATen/ops/reflection_pad2d_backward_cuda_dispatch.h>
|
| 494 |
+
#include <ATen/ops/reflection_pad3d_cuda_dispatch.h>
|
| 495 |
+
#include <ATen/ops/reflection_pad3d_backward_cuda_dispatch.h>
|
| 496 |
+
#include <ATen/ops/relu_cuda_dispatch.h>
|
| 497 |
+
#include <ATen/ops/remainder_cuda_dispatch.h>
|
| 498 |
+
#include <ATen/ops/renorm_cuda_dispatch.h>
|
| 499 |
+
#include <ATen/ops/repeat_interleave_cuda_dispatch.h>
|
| 500 |
+
#include <ATen/ops/replication_pad1d_cuda_dispatch.h>
|
| 501 |
+
#include <ATen/ops/replication_pad1d_backward_cuda_dispatch.h>
|
| 502 |
+
#include <ATen/ops/replication_pad2d_cuda_dispatch.h>
|
| 503 |
+
#include <ATen/ops/replication_pad2d_backward_cuda_dispatch.h>
|
| 504 |
+
#include <ATen/ops/replication_pad3d_cuda_dispatch.h>
|
| 505 |
+
#include <ATen/ops/replication_pad3d_backward_cuda_dispatch.h>
|
| 506 |
+
#include <ATen/ops/resize_cuda_dispatch.h>
|
| 507 |
+
#include <ATen/ops/roll_cuda_dispatch.h>
|
| 508 |
+
#include <ATen/ops/round_cuda_dispatch.h>
|
| 509 |
+
#include <ATen/ops/rrelu_with_noise_cuda_dispatch.h>
|
| 510 |
+
#include <ATen/ops/rshift_cuda_dispatch.h>
|
| 511 |
+
#include <ATen/ops/rsqrt_cuda_dispatch.h>
|
| 512 |
+
#include <ATen/ops/rsub_cuda_dispatch.h>
|
| 513 |
+
#include <ATen/ops/scatter_cuda_dispatch.h>
|
| 514 |
+
#include <ATen/ops/scatter_add_cuda_dispatch.h>
|
| 515 |
+
#include <ATen/ops/scatter_reduce_cuda_dispatch.h>
|
| 516 |
+
#include <ATen/ops/searchsorted_cuda_dispatch.h>
|
| 517 |
+
#include <ATen/ops/segment_reduce_cuda_dispatch.h>
|
| 518 |
+
#include <ATen/ops/set_cuda_dispatch.h>
|
| 519 |
+
#include <ATen/ops/sgn_cuda_dispatch.h>
|
| 520 |
+
#include <ATen/ops/sigmoid_cuda_dispatch.h>
|
| 521 |
+
#include <ATen/ops/sigmoid_backward_cuda_dispatch.h>
|
| 522 |
+
#include <ATen/ops/sign_cuda_dispatch.h>
|
| 523 |
+
#include <ATen/ops/signbit_cuda_dispatch.h>
|
| 524 |
+
#include <ATen/ops/silu_cuda_dispatch.h>
|
| 525 |
+
#include <ATen/ops/silu_backward_cuda_dispatch.h>
|
| 526 |
+
#include <ATen/ops/sin_cuda_dispatch.h>
|
| 527 |
+
#include <ATen/ops/sinc_cuda_dispatch.h>
|
| 528 |
+
#include <ATen/ops/sinh_cuda_dispatch.h>
|
| 529 |
+
#include <ATen/ops/slow_conv_dilated2d_cuda_dispatch.h>
|
| 530 |
+
#include <ATen/ops/slow_conv_dilated3d_cuda_dispatch.h>
|
| 531 |
+
#include <ATen/ops/slow_conv_transpose2d_cuda_dispatch.h>
|
| 532 |
+
#include <ATen/ops/slow_conv_transpose3d_cuda_dispatch.h>
|
| 533 |
+
#include <ATen/ops/smooth_l1_loss_cuda_dispatch.h>
|
| 534 |
+
#include <ATen/ops/smooth_l1_loss_backward_cuda_dispatch.h>
|
| 535 |
+
#include <ATen/ops/softplus_cuda_dispatch.h>
|
| 536 |
+
#include <ATen/ops/softplus_backward_cuda_dispatch.h>
|
| 537 |
+
#include <ATen/ops/softshrink_cuda_dispatch.h>
|
| 538 |
+
#include <ATen/ops/softshrink_backward_cuda_dispatch.h>
|
| 539 |
+
#include <ATen/ops/sort_cuda_dispatch.h>
|
| 540 |
+
#include <ATen/ops/special_airy_ai_cuda_dispatch.h>
|
| 541 |
+
#include <ATen/ops/special_bessel_j0_cuda_dispatch.h>
|
| 542 |
+
#include <ATen/ops/special_bessel_j1_cuda_dispatch.h>
|
| 543 |
+
#include <ATen/ops/special_bessel_y0_cuda_dispatch.h>
|
| 544 |
+
#include <ATen/ops/special_bessel_y1_cuda_dispatch.h>
|
| 545 |
+
#include <ATen/ops/special_chebyshev_polynomial_t_cuda_dispatch.h>
|
| 546 |
+
#include <ATen/ops/special_chebyshev_polynomial_u_cuda_dispatch.h>
|
| 547 |
+
#include <ATen/ops/special_chebyshev_polynomial_v_cuda_dispatch.h>
|
| 548 |
+
#include <ATen/ops/special_chebyshev_polynomial_w_cuda_dispatch.h>
|
| 549 |
+
#include <ATen/ops/special_entr_cuda_dispatch.h>
|
| 550 |
+
#include <ATen/ops/special_erfcx_cuda_dispatch.h>
|
| 551 |
+
#include <ATen/ops/special_hermite_polynomial_h_cuda_dispatch.h>
|
| 552 |
+
#include <ATen/ops/special_hermite_polynomial_he_cuda_dispatch.h>
|
| 553 |
+
#include <ATen/ops/special_i0e_cuda_dispatch.h>
|
| 554 |
+
#include <ATen/ops/special_i1_cuda_dispatch.h>
|
| 555 |
+
#include <ATen/ops/special_i1e_cuda_dispatch.h>
|
| 556 |
+
#include <ATen/ops/special_laguerre_polynomial_l_cuda_dispatch.h>
|
| 557 |
+
#include <ATen/ops/special_legendre_polynomial_p_cuda_dispatch.h>
|
| 558 |
+
#include <ATen/ops/special_log_ndtr_cuda_dispatch.h>
|
| 559 |
+
#include <ATen/ops/special_modified_bessel_i0_cuda_dispatch.h>
|
| 560 |
+
#include <ATen/ops/special_modified_bessel_i1_cuda_dispatch.h>
|
| 561 |
+
#include <ATen/ops/special_modified_bessel_k0_cuda_dispatch.h>
|
| 562 |
+
#include <ATen/ops/special_modified_bessel_k1_cuda_dispatch.h>
|
| 563 |
+
#include <ATen/ops/special_ndtri_cuda_dispatch.h>
|
| 564 |
+
#include <ATen/ops/special_scaled_modified_bessel_k0_cuda_dispatch.h>
|
| 565 |
+
#include <ATen/ops/special_scaled_modified_bessel_k1_cuda_dispatch.h>
|
| 566 |
+
#include <ATen/ops/special_shifted_chebyshev_polynomial_t_cuda_dispatch.h>
|
| 567 |
+
#include <ATen/ops/special_shifted_chebyshev_polynomial_u_cuda_dispatch.h>
|
| 568 |
+
#include <ATen/ops/special_shifted_chebyshev_polynomial_v_cuda_dispatch.h>
|
| 569 |
+
#include <ATen/ops/special_shifted_chebyshev_polynomial_w_cuda_dispatch.h>
|
| 570 |
+
#include <ATen/ops/special_spherical_bessel_j0_cuda_dispatch.h>
|
| 571 |
+
#include <ATen/ops/special_xlog1py_cuda_dispatch.h>
|
| 572 |
+
#include <ATen/ops/special_zeta_cuda_dispatch.h>
|
| 573 |
+
#include <ATen/ops/split_with_sizes_copy_cuda_dispatch.h>
|
| 574 |
+
#include <ATen/ops/sqrt_cuda_dispatch.h>
|
| 575 |
+
#include <ATen/ops/sspaddmm_cuda_dispatch.h>
|
| 576 |
+
#include <ATen/ops/std_cuda_dispatch.h>
|
| 577 |
+
#include <ATen/ops/std_mean_cuda_dispatch.h>
|
| 578 |
+
#include <ATen/ops/sub_cuda_dispatch.h>
|
| 579 |
+
#include <ATen/ops/sum_cuda_dispatch.h>
|
| 580 |
+
#include <ATen/ops/take_cuda_dispatch.h>
|
| 581 |
+
#include <ATen/ops/tan_cuda_dispatch.h>
|
| 582 |
+
#include <ATen/ops/tanh_cuda_dispatch.h>
|
| 583 |
+
#include <ATen/ops/tanh_backward_cuda_dispatch.h>
|
| 584 |
+
#include <ATen/ops/threshold_cuda_dispatch.h>
|
| 585 |
+
#include <ATen/ops/threshold_backward_cuda_dispatch.h>
|
| 586 |
+
#include <ATen/ops/topk_cuda_dispatch.h>
|
| 587 |
+
#include <ATen/ops/trace_cuda_dispatch.h>
|
| 588 |
+
#include <ATen/ops/triangular_solve_cuda_dispatch.h>
|
| 589 |
+
#include <ATen/ops/tril_cuda_dispatch.h>
|
| 590 |
+
#include <ATen/ops/tril_indices_cuda_dispatch.h>
|
| 591 |
+
#include <ATen/ops/triu_cuda_dispatch.h>
|
| 592 |
+
#include <ATen/ops/triu_indices_cuda_dispatch.h>
|
| 593 |
+
#include <ATen/ops/trunc_cuda_dispatch.h>
|
| 594 |
+
#include <ATen/ops/unfold_cuda_dispatch.h>
|
| 595 |
+
#include <ATen/ops/unfold_backward_cuda_dispatch.h>
|
| 596 |
+
#include <ATen/ops/uniform_cuda_dispatch.h>
|
| 597 |
+
#include <ATen/ops/unique_consecutive_cuda_dispatch.h>
|
| 598 |
+
#include <ATen/ops/unique_dim_cuda_dispatch.h>
|
| 599 |
+
#include <ATen/ops/unique_dim_consecutive_cuda_dispatch.h>
|
| 600 |
+
#include <ATen/ops/upsample_bicubic2d_cuda_dispatch.h>
|
| 601 |
+
#include <ATen/ops/upsample_bicubic2d_backward_cuda_dispatch.h>
|
| 602 |
+
#include <ATen/ops/upsample_bilinear2d_cuda_dispatch.h>
|
| 603 |
+
#include <ATen/ops/upsample_bilinear2d_backward_cuda_dispatch.h>
|
| 604 |
+
#include <ATen/ops/upsample_linear1d_cuda_dispatch.h>
|
| 605 |
+
#include <ATen/ops/upsample_linear1d_backward_cuda_dispatch.h>
|
| 606 |
+
#include <ATen/ops/upsample_nearest1d_cuda_dispatch.h>
|
| 607 |
+
#include <ATen/ops/upsample_nearest1d_backward_cuda_dispatch.h>
|
| 608 |
+
#include <ATen/ops/upsample_nearest2d_cuda_dispatch.h>
|
| 609 |
+
#include <ATen/ops/upsample_nearest2d_backward_cuda_dispatch.h>
|
| 610 |
+
#include <ATen/ops/upsample_nearest3d_cuda_dispatch.h>
|
| 611 |
+
#include <ATen/ops/upsample_nearest3d_backward_cuda_dispatch.h>
|
| 612 |
+
#include <ATen/ops/upsample_trilinear3d_cuda_dispatch.h>
|
| 613 |
+
#include <ATen/ops/upsample_trilinear3d_backward_cuda_dispatch.h>
|
| 614 |
+
#include <ATen/ops/var_cuda_dispatch.h>
|
| 615 |
+
#include <ATen/ops/var_mean_cuda_dispatch.h>
|
| 616 |
+
#include <ATen/ops/vdot_cuda_dispatch.h>
|
| 617 |
+
#include <ATen/ops/view_cuda_dispatch.h>
|
| 618 |
+
#include <ATen/ops/view_as_complex_cuda_dispatch.h>
|
| 619 |
+
#include <ATen/ops/view_as_real_cuda_dispatch.h>
|
| 620 |
+
#include <ATen/ops/where_cuda_dispatch.h>
|
| 621 |
+
#include <ATen/ops/xlogy_cuda_dispatch.h>
|
| 622 |
+
#include <ATen/ops/zero_cuda_dispatch.h>
|
| 623 |
+
|
| 624 |
+
|
| 625 |
+
|
lib/python3.12/site-packages/torch/include/ATen/CachedTensorUtils.h
ADDED
|
@@ -0,0 +1,24 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <ATen/ATen.h>
|
| 4 |
+
|
| 5 |
+
namespace at::caching {
|
| 6 |
+
|
| 7 |
+
// Some systems (just cudagraphs currently) will persist a static tensor output
|
| 8 |
+
// whose TensorImpl does not change across iterations. For these tensors caching
|
| 9 |
+
// dtype conversions is invalid. Additionally, there will be an extra reference
|
| 10 |
+
// count to these cached tensors that would prevent buffer inplacing and other
|
| 11 |
+
// checks on tensor uniqueness. If we are not using these systems the enabled
|
| 12 |
+
// flag will be false and we will avoid the hash lookup.
|
| 13 |
+
|
| 14 |
+
TORCH_API bool is_cached_tensor(const at::Tensor& t);
|
| 15 |
+
TORCH_API void add_cached_tensor(const at::Tensor& t);
|
| 16 |
+
TORCH_API void remove_cached_tensor(const at::Tensor& t);
|
| 17 |
+
TORCH_API void set_cached_tensors_enabled(bool enable);
|
| 18 |
+
|
| 19 |
+
// For gradient buffer stealing we will adjust the use count of tensors
|
| 20 |
+
// which are persisted by cudagraphs, just as we need to adjust reference
|
| 21 |
+
// count of tensors with hooks.
|
| 22 |
+
TORCH_API size_t adjusted_use_count(const at::Tensor& t);
|
| 23 |
+
|
| 24 |
+
} // namespace at::caching
|
lib/python3.12/site-packages/torch/include/ATen/CompositeExplicitAutogradFunctions_inl.h
ADDED
|
@@ -0,0 +1,560 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
// @generated by torchgen/gen.py from DispatchKeyFunctions_inl.h
|
| 3 |
+
|
| 4 |
+
// NB: The implementing C++ file is RegisterDispatchKey.cpp
|
| 5 |
+
|
| 6 |
+
// The only #includes we need are for custom classes that have defaults in the C++ API
|
| 7 |
+
#include <c10/core/MemoryFormat.h>
|
| 8 |
+
#include <c10/core/Scalar.h>
|
| 9 |
+
#include <ATen/core/Reduction.h>
|
| 10 |
+
|
| 11 |
+
#if defined(AT_PER_OPERATOR_HEADERS) && defined(TORCH_ASSERT_ONLY_METHOD_OPERATORS)
|
| 12 |
+
#error This change adds a dependency on all pytorch operators, meaning the \
|
| 13 |
+
file will need to be re-compiled every time an operator is changed or added. \
|
| 14 |
+
Consider including a specific operator from \
|
| 15 |
+
<ATen/ops/{my_operator}_compositeexplicitautograd_dispatch.h>. \
|
| 16 |
+
See NOTE [TORCH_ASSERT_ONLY_METHOD_OPERATORS].
|
| 17 |
+
#endif
|
| 18 |
+
|
| 19 |
+
#include <ATen/ops/_adaptive_avg_pool2d_compositeexplicitautograd_dispatch.h>
|
| 20 |
+
#include <ATen/ops/_adaptive_avg_pool2d_backward_compositeexplicitautograd_dispatch.h>
|
| 21 |
+
#include <ATen/ops/_adaptive_avg_pool3d_compositeexplicitautograd_dispatch.h>
|
| 22 |
+
#include <ATen/ops/_adaptive_avg_pool3d_backward_compositeexplicitautograd_dispatch.h>
|
| 23 |
+
#include <ATen/ops/_add_relu_compositeexplicitautograd_dispatch.h>
|
| 24 |
+
#include <ATen/ops/_aminmax_compositeexplicitautograd_dispatch.h>
|
| 25 |
+
#include <ATen/ops/_amp_foreach_non_finite_check_and_unscale_compositeexplicitautograd_dispatch.h>
|
| 26 |
+
#include <ATen/ops/_amp_update_scale_compositeexplicitautograd_dispatch.h>
|
| 27 |
+
#include <ATen/ops/_assert_scalar_compositeexplicitautograd_dispatch.h>
|
| 28 |
+
#include <ATen/ops/_assert_tensor_metadata_compositeexplicitautograd_dispatch.h>
|
| 29 |
+
#include <ATen/ops/_batch_norm_no_update_compositeexplicitautograd_dispatch.h>
|
| 30 |
+
#include <ATen/ops/_batch_norm_with_update_compositeexplicitautograd_dispatch.h>
|
| 31 |
+
#include <ATen/ops/_cdist_backward_compositeexplicitautograd_dispatch.h>
|
| 32 |
+
#include <ATen/ops/_cdist_forward_compositeexplicitautograd_dispatch.h>
|
| 33 |
+
#include <ATen/ops/_cholesky_solve_helper_compositeexplicitautograd_dispatch.h>
|
| 34 |
+
#include <ATen/ops/_chunk_cat_compositeexplicitautograd_dispatch.h>
|
| 35 |
+
#include <ATen/ops/_coalesce_compositeexplicitautograd_dispatch.h>
|
| 36 |
+
#include <ATen/ops/_coalesced_compositeexplicitautograd_dispatch.h>
|
| 37 |
+
#include <ATen/ops/_conj_compositeexplicitautograd_dispatch.h>
|
| 38 |
+
#include <ATen/ops/_conj_copy_compositeexplicitautograd_dispatch.h>
|
| 39 |
+
#include <ATen/ops/_conj_physical_compositeexplicitautograd_dispatch.h>
|
| 40 |
+
#include <ATen/ops/_convolution_compositeexplicitautograd_dispatch.h>
|
| 41 |
+
#include <ATen/ops/_copy_from_compositeexplicitautograd_dispatch.h>
|
| 42 |
+
#include <ATen/ops/_copy_from_and_resize_compositeexplicitautograd_dispatch.h>
|
| 43 |
+
#include <ATen/ops/_ctc_loss_compositeexplicitautograd_dispatch.h>
|
| 44 |
+
#include <ATen/ops/_ctc_loss_backward_compositeexplicitautograd_dispatch.h>
|
| 45 |
+
#include <ATen/ops/_cudnn_ctc_loss_compositeexplicitautograd_dispatch.h>
|
| 46 |
+
#include <ATen/ops/_cudnn_init_dropout_state_compositeexplicitautograd_dispatch.h>
|
| 47 |
+
#include <ATen/ops/_cudnn_rnn_compositeexplicitautograd_dispatch.h>
|
| 48 |
+
#include <ATen/ops/_cudnn_rnn_backward_compositeexplicitautograd_dispatch.h>
|
| 49 |
+
#include <ATen/ops/_cudnn_rnn_flatten_weight_compositeexplicitautograd_dispatch.h>
|
| 50 |
+
#include <ATen/ops/_dirichlet_grad_compositeexplicitautograd_dispatch.h>
|
| 51 |
+
#include <ATen/ops/_efficientzerotensor_compositeexplicitautograd_dispatch.h>
|
| 52 |
+
#include <ATen/ops/_embedding_bag_compositeexplicitautograd_dispatch.h>
|
| 53 |
+
#include <ATen/ops/_embedding_bag_dense_backward_compositeexplicitautograd_dispatch.h>
|
| 54 |
+
#include <ATen/ops/_embedding_bag_forward_only_compositeexplicitautograd_dispatch.h>
|
| 55 |
+
#include <ATen/ops/_embedding_bag_per_sample_weights_backward_compositeexplicitautograd_dispatch.h>
|
| 56 |
+
#include <ATen/ops/_empty_affine_quantized_compositeexplicitautograd_dispatch.h>
|
| 57 |
+
#include <ATen/ops/_empty_per_channel_affine_quantized_compositeexplicitautograd_dispatch.h>
|
| 58 |
+
#include <ATen/ops/_euclidean_dist_compositeexplicitautograd_dispatch.h>
|
| 59 |
+
#include <ATen/ops/_fake_quantize_learnable_per_channel_affine_compositeexplicitautograd_dispatch.h>
|
| 60 |
+
#include <ATen/ops/_fake_quantize_learnable_per_tensor_affine_compositeexplicitautograd_dispatch.h>
|
| 61 |
+
#include <ATen/ops/_fake_quantize_per_tensor_affine_cachemask_tensor_qparams_compositeexplicitautograd_dispatch.h>
|
| 62 |
+
#include <ATen/ops/_foobar_compositeexplicitautograd_dispatch.h>
|
| 63 |
+
#include <ATen/ops/_foreach_abs_compositeexplicitautograd_dispatch.h>
|
| 64 |
+
#include <ATen/ops/_foreach_acos_compositeexplicitautograd_dispatch.h>
|
| 65 |
+
#include <ATen/ops/_foreach_add_compositeexplicitautograd_dispatch.h>
|
| 66 |
+
#include <ATen/ops/_foreach_addcdiv_compositeexplicitautograd_dispatch.h>
|
| 67 |
+
#include <ATen/ops/_foreach_addcmul_compositeexplicitautograd_dispatch.h>
|
| 68 |
+
#include <ATen/ops/_foreach_asin_compositeexplicitautograd_dispatch.h>
|
| 69 |
+
#include <ATen/ops/_foreach_atan_compositeexplicitautograd_dispatch.h>
|
| 70 |
+
#include <ATen/ops/_foreach_ceil_compositeexplicitautograd_dispatch.h>
|
| 71 |
+
#include <ATen/ops/_foreach_clamp_max_compositeexplicitautograd_dispatch.h>
|
| 72 |
+
#include <ATen/ops/_foreach_clamp_min_compositeexplicitautograd_dispatch.h>
|
| 73 |
+
#include <ATen/ops/_foreach_copy_compositeexplicitautograd_dispatch.h>
|
| 74 |
+
#include <ATen/ops/_foreach_cos_compositeexplicitautograd_dispatch.h>
|
| 75 |
+
#include <ATen/ops/_foreach_cosh_compositeexplicitautograd_dispatch.h>
|
| 76 |
+
#include <ATen/ops/_foreach_div_compositeexplicitautograd_dispatch.h>
|
| 77 |
+
#include <ATen/ops/_foreach_erf_compositeexplicitautograd_dispatch.h>
|
| 78 |
+
#include <ATen/ops/_foreach_erfc_compositeexplicitautograd_dispatch.h>
|
| 79 |
+
#include <ATen/ops/_foreach_exp_compositeexplicitautograd_dispatch.h>
|
| 80 |
+
#include <ATen/ops/_foreach_expm1_compositeexplicitautograd_dispatch.h>
|
| 81 |
+
#include <ATen/ops/_foreach_floor_compositeexplicitautograd_dispatch.h>
|
| 82 |
+
#include <ATen/ops/_foreach_frac_compositeexplicitautograd_dispatch.h>
|
| 83 |
+
#include <ATen/ops/_foreach_lerp_compositeexplicitautograd_dispatch.h>
|
| 84 |
+
#include <ATen/ops/_foreach_lgamma_compositeexplicitautograd_dispatch.h>
|
| 85 |
+
#include <ATen/ops/_foreach_log_compositeexplicitautograd_dispatch.h>
|
| 86 |
+
#include <ATen/ops/_foreach_log10_compositeexplicitautograd_dispatch.h>
|
| 87 |
+
#include <ATen/ops/_foreach_log1p_compositeexplicitautograd_dispatch.h>
|
| 88 |
+
#include <ATen/ops/_foreach_log2_compositeexplicitautograd_dispatch.h>
|
| 89 |
+
#include <ATen/ops/_foreach_max_compositeexplicitautograd_dispatch.h>
|
| 90 |
+
#include <ATen/ops/_foreach_maximum_compositeexplicitautograd_dispatch.h>
|
| 91 |
+
#include <ATen/ops/_foreach_minimum_compositeexplicitautograd_dispatch.h>
|
| 92 |
+
#include <ATen/ops/_foreach_mul_compositeexplicitautograd_dispatch.h>
|
| 93 |
+
#include <ATen/ops/_foreach_neg_compositeexplicitautograd_dispatch.h>
|
| 94 |
+
#include <ATen/ops/_foreach_norm_compositeexplicitautograd_dispatch.h>
|
| 95 |
+
#include <ATen/ops/_foreach_pow_compositeexplicitautograd_dispatch.h>
|
| 96 |
+
#include <ATen/ops/_foreach_reciprocal_compositeexplicitautograd_dispatch.h>
|
| 97 |
+
#include <ATen/ops/_foreach_round_compositeexplicitautograd_dispatch.h>
|
| 98 |
+
#include <ATen/ops/_foreach_rsqrt_compositeexplicitautograd_dispatch.h>
|
| 99 |
+
#include <ATen/ops/_foreach_sigmoid_compositeexplicitautograd_dispatch.h>
|
| 100 |
+
#include <ATen/ops/_foreach_sign_compositeexplicitautograd_dispatch.h>
|
| 101 |
+
#include <ATen/ops/_foreach_sin_compositeexplicitautograd_dispatch.h>
|
| 102 |
+
#include <ATen/ops/_foreach_sinh_compositeexplicitautograd_dispatch.h>
|
| 103 |
+
#include <ATen/ops/_foreach_sqrt_compositeexplicitautograd_dispatch.h>
|
| 104 |
+
#include <ATen/ops/_foreach_sub_compositeexplicitautograd_dispatch.h>
|
| 105 |
+
#include <ATen/ops/_foreach_tan_compositeexplicitautograd_dispatch.h>
|
| 106 |
+
#include <ATen/ops/_foreach_tanh_compositeexplicitautograd_dispatch.h>
|
| 107 |
+
#include <ATen/ops/_foreach_trunc_compositeexplicitautograd_dispatch.h>
|
| 108 |
+
#include <ATen/ops/_foreach_zero_compositeexplicitautograd_dispatch.h>
|
| 109 |
+
#include <ATen/ops/_functional_assert_scalar_compositeexplicitautograd_dispatch.h>
|
| 110 |
+
#include <ATen/ops/_functional_sym_constrain_range_compositeexplicitautograd_dispatch.h>
|
| 111 |
+
#include <ATen/ops/_functional_sym_constrain_range_for_size_compositeexplicitautograd_dispatch.h>
|
| 112 |
+
#include <ATen/ops/_fused_adagrad_compositeexplicitautograd_dispatch.h>
|
| 113 |
+
#include <ATen/ops/_fused_adam_compositeexplicitautograd_dispatch.h>
|
| 114 |
+
#include <ATen/ops/_fused_adamw_compositeexplicitautograd_dispatch.h>
|
| 115 |
+
#include <ATen/ops/_fused_dropout_compositeexplicitautograd_dispatch.h>
|
| 116 |
+
#include <ATen/ops/_fused_moving_avg_obs_fq_helper_compositeexplicitautograd_dispatch.h>
|
| 117 |
+
#include <ATen/ops/_fused_sgd_compositeexplicitautograd_dispatch.h>
|
| 118 |
+
#include <ATen/ops/_fw_primal_compositeexplicitautograd_dispatch.h>
|
| 119 |
+
#include <ATen/ops/_fw_primal_copy_compositeexplicitautograd_dispatch.h>
|
| 120 |
+
#include <ATen/ops/_grid_sampler_2d_cpu_fallback_compositeexplicitautograd_dispatch.h>
|
| 121 |
+
#include <ATen/ops/_has_same_storage_numel_compositeexplicitautograd_dispatch.h>
|
| 122 |
+
#include <ATen/ops/_histogramdd_bin_edges_compositeexplicitautograd_dispatch.h>
|
| 123 |
+
#include <ATen/ops/_histogramdd_from_bin_cts_compositeexplicitautograd_dispatch.h>
|
| 124 |
+
#include <ATen/ops/_histogramdd_from_bin_tensors_compositeexplicitautograd_dispatch.h>
|
| 125 |
+
#include <ATen/ops/_index_put_impl_compositeexplicitautograd_dispatch.h>
|
| 126 |
+
#include <ATen/ops/_indices_copy_compositeexplicitautograd_dispatch.h>
|
| 127 |
+
#include <ATen/ops/_is_all_true_compositeexplicitautograd_dispatch.h>
|
| 128 |
+
#include <ATen/ops/_is_any_true_compositeexplicitautograd_dispatch.h>
|
| 129 |
+
#include <ATen/ops/_lazy_clone_compositeexplicitautograd_dispatch.h>
|
| 130 |
+
#include <ATen/ops/_linalg_check_errors_compositeexplicitautograd_dispatch.h>
|
| 131 |
+
#include <ATen/ops/_lstm_mps_compositeexplicitautograd_dispatch.h>
|
| 132 |
+
#include <ATen/ops/_make_dual_compositeexplicitautograd_dispatch.h>
|
| 133 |
+
#include <ATen/ops/_make_dual_copy_compositeexplicitautograd_dispatch.h>
|
| 134 |
+
#include <ATen/ops/_make_per_channel_quantized_tensor_compositeexplicitautograd_dispatch.h>
|
| 135 |
+
#include <ATen/ops/_make_per_tensor_quantized_tensor_compositeexplicitautograd_dispatch.h>
|
| 136 |
+
#include <ATen/ops/_masked_scale_compositeexplicitautograd_dispatch.h>
|
| 137 |
+
#include <ATen/ops/_masked_softmax_compositeexplicitautograd_dispatch.h>
|
| 138 |
+
#include <ATen/ops/_masked_softmax_backward_compositeexplicitautograd_dispatch.h>
|
| 139 |
+
#include <ATen/ops/_mkldnn_reshape_compositeexplicitautograd_dispatch.h>
|
| 140 |
+
#include <ATen/ops/_mkldnn_transpose_compositeexplicitautograd_dispatch.h>
|
| 141 |
+
#include <ATen/ops/_mps_convolution_compositeexplicitautograd_dispatch.h>
|
| 142 |
+
#include <ATen/ops/_mps_convolution_transpose_compositeexplicitautograd_dispatch.h>
|
| 143 |
+
#include <ATen/ops/_native_batch_norm_legit_compositeexplicitautograd_dispatch.h>
|
| 144 |
+
#include <ATen/ops/_native_batch_norm_legit_no_training_compositeexplicitautograd_dispatch.h>
|
| 145 |
+
#include <ATen/ops/_native_multi_head_attention_compositeexplicitautograd_dispatch.h>
|
| 146 |
+
#include <ATen/ops/_neg_view_compositeexplicitautograd_dispatch.h>
|
| 147 |
+
#include <ATen/ops/_neg_view_copy_compositeexplicitautograd_dispatch.h>
|
| 148 |
+
#include <ATen/ops/_nested_from_padded_compositeexplicitautograd_dispatch.h>
|
| 149 |
+
#include <ATen/ops/_nested_from_padded_and_nested_example_compositeexplicitautograd_dispatch.h>
|
| 150 |
+
#include <ATen/ops/_nested_get_values_copy_compositeexplicitautograd_dispatch.h>
|
| 151 |
+
#include <ATen/ops/_nested_tensor_from_mask_compositeexplicitautograd_dispatch.h>
|
| 152 |
+
#include <ATen/ops/_nested_tensor_from_tensor_list_compositeexplicitautograd_dispatch.h>
|
| 153 |
+
#include <ATen/ops/_nested_tensor_size_compositeexplicitautograd_dispatch.h>
|
| 154 |
+
#include <ATen/ops/_nested_tensor_storage_offsets_compositeexplicitautograd_dispatch.h>
|
| 155 |
+
#include <ATen/ops/_nested_tensor_strides_compositeexplicitautograd_dispatch.h>
|
| 156 |
+
#include <ATen/ops/_nested_view_from_buffer_copy_compositeexplicitautograd_dispatch.h>
|
| 157 |
+
#include <ATen/ops/_nested_view_from_jagged_copy_compositeexplicitautograd_dispatch.h>
|
| 158 |
+
#include <ATen/ops/_new_zeros_with_same_feature_meta_compositeexplicitautograd_dispatch.h>
|
| 159 |
+
#include <ATen/ops/_nnpack_spatial_convolution_compositeexplicitautograd_dispatch.h>
|
| 160 |
+
#include <ATen/ops/_pack_padded_sequence_compositeexplicitautograd_dispatch.h>
|
| 161 |
+
#include <ATen/ops/_pdist_backward_compositeexplicitautograd_dispatch.h>
|
| 162 |
+
#include <ATen/ops/_pdist_forward_compositeexplicitautograd_dispatch.h>
|
| 163 |
+
#include <ATen/ops/_pin_memory_compositeexplicitautograd_dispatch.h>
|
| 164 |
+
#include <ATen/ops/_print_compositeexplicitautograd_dispatch.h>
|
| 165 |
+
#include <ATen/ops/_reshape_alias_copy_compositeexplicitautograd_dispatch.h>
|
| 166 |
+
#include <ATen/ops/_reshape_copy_compositeexplicitautograd_dispatch.h>
|
| 167 |
+
#include <ATen/ops/_resize_output_compositeexplicitautograd_dispatch.h>
|
| 168 |
+
#include <ATen/ops/_safe_softmax_compositeexplicitautograd_dispatch.h>
|
| 169 |
+
#include <ATen/ops/_sample_dirichlet_compositeexplicitautograd_dispatch.h>
|
| 170 |
+
#include <ATen/ops/_scaled_dot_product_fused_attention_overrideable_compositeexplicitautograd_dispatch.h>
|
| 171 |
+
#include <ATen/ops/_scaled_dot_product_fused_attention_overrideable_backward_compositeexplicitautograd_dispatch.h>
|
| 172 |
+
#include <ATen/ops/_segment_reduce_backward_compositeexplicitautograd_dispatch.h>
|
| 173 |
+
#include <ATen/ops/_slow_conv2d_backward_compositeexplicitautograd_dispatch.h>
|
| 174 |
+
#include <ATen/ops/_sparse_addmm_compositeexplicitautograd_dispatch.h>
|
| 175 |
+
#include <ATen/ops/_sparse_broadcast_to_copy_compositeexplicitautograd_dispatch.h>
|
| 176 |
+
#include <ATen/ops/_sparse_compressed_tensor_with_dims_compositeexplicitautograd_dispatch.h>
|
| 177 |
+
#include <ATen/ops/_sparse_coo_tensor_with_dims_compositeexplicitautograd_dispatch.h>
|
| 178 |
+
#include <ATen/ops/_sparse_coo_tensor_with_dims_and_tensors_compositeexplicitautograd_dispatch.h>
|
| 179 |
+
#include <ATen/ops/_sparse_csr_prod_compositeexplicitautograd_dispatch.h>
|
| 180 |
+
#include <ATen/ops/_sparse_csr_sum_compositeexplicitautograd_dispatch.h>
|
| 181 |
+
#include <ATen/ops/_sparse_log_softmax_compositeexplicitautograd_dispatch.h>
|
| 182 |
+
#include <ATen/ops/_sparse_log_softmax_backward_data_compositeexplicitautograd_dispatch.h>
|
| 183 |
+
#include <ATen/ops/_sparse_mask_projection_compositeexplicitautograd_dispatch.h>
|
| 184 |
+
#include <ATen/ops/_sparse_softmax_compositeexplicitautograd_dispatch.h>
|
| 185 |
+
#include <ATen/ops/_sparse_softmax_backward_data_compositeexplicitautograd_dispatch.h>
|
| 186 |
+
#include <ATen/ops/_sparse_sparse_matmul_compositeexplicitautograd_dispatch.h>
|
| 187 |
+
#include <ATen/ops/_sparse_sum_compositeexplicitautograd_dispatch.h>
|
| 188 |
+
#include <ATen/ops/_sparse_sum_backward_compositeexplicitautograd_dispatch.h>
|
| 189 |
+
#include <ATen/ops/_spdiags_compositeexplicitautograd_dispatch.h>
|
| 190 |
+
#include <ATen/ops/_stack_compositeexplicitautograd_dispatch.h>
|
| 191 |
+
#include <ATen/ops/_standard_gamma_compositeexplicitautograd_dispatch.h>
|
| 192 |
+
#include <ATen/ops/_standard_gamma_grad_compositeexplicitautograd_dispatch.h>
|
| 193 |
+
#include <ATen/ops/_test_autograd_multiple_dispatch_compositeexplicitautograd_dispatch.h>
|
| 194 |
+
#include <ATen/ops/_test_autograd_multiple_dispatch_view_compositeexplicitautograd_dispatch.h>
|
| 195 |
+
#include <ATen/ops/_test_autograd_multiple_dispatch_view_copy_compositeexplicitautograd_dispatch.h>
|
| 196 |
+
#include <ATen/ops/_test_functorch_fallback_compositeexplicitautograd_dispatch.h>
|
| 197 |
+
#include <ATen/ops/_test_optional_filled_intlist_compositeexplicitautograd_dispatch.h>
|
| 198 |
+
#include <ATen/ops/_test_optional_floatlist_compositeexplicitautograd_dispatch.h>
|
| 199 |
+
#include <ATen/ops/_test_optional_intlist_compositeexplicitautograd_dispatch.h>
|
| 200 |
+
#include <ATen/ops/_test_parallel_materialize_compositeexplicitautograd_dispatch.h>
|
| 201 |
+
#include <ATen/ops/_test_warn_in_autograd_compositeexplicitautograd_dispatch.h>
|
| 202 |
+
#include <ATen/ops/_thnn_fused_gru_cell_compositeexplicitautograd_dispatch.h>
|
| 203 |
+
#include <ATen/ops/_thnn_fused_gru_cell_backward_compositeexplicitautograd_dispatch.h>
|
| 204 |
+
#include <ATen/ops/_thnn_fused_lstm_cell_compositeexplicitautograd_dispatch.h>
|
| 205 |
+
#include <ATen/ops/_thnn_fused_lstm_cell_backward_impl_compositeexplicitautograd_dispatch.h>
|
| 206 |
+
#include <ATen/ops/_to_copy_compositeexplicitautograd_dispatch.h>
|
| 207 |
+
#include <ATen/ops/_to_dense_compositeexplicitautograd_dispatch.h>
|
| 208 |
+
#include <ATen/ops/_to_sparse_compositeexplicitautograd_dispatch.h>
|
| 209 |
+
#include <ATen/ops/_to_sparse_bsc_compositeexplicitautograd_dispatch.h>
|
| 210 |
+
#include <ATen/ops/_to_sparse_bsr_compositeexplicitautograd_dispatch.h>
|
| 211 |
+
#include <ATen/ops/_to_sparse_csc_compositeexplicitautograd_dispatch.h>
|
| 212 |
+
#include <ATen/ops/_to_sparse_csr_compositeexplicitautograd_dispatch.h>
|
| 213 |
+
#include <ATen/ops/_transform_bias_rescale_qkv_compositeexplicitautograd_dispatch.h>
|
| 214 |
+
#include <ATen/ops/_transformer_encoder_layer_fwd_compositeexplicitautograd_dispatch.h>
|
| 215 |
+
#include <ATen/ops/_trilinear_compositeexplicitautograd_dispatch.h>
|
| 216 |
+
#include <ATen/ops/_triton_multi_head_attention_compositeexplicitautograd_dispatch.h>
|
| 217 |
+
#include <ATen/ops/_triton_scaled_dot_attention_compositeexplicitautograd_dispatch.h>
|
| 218 |
+
#include <ATen/ops/_unique_compositeexplicitautograd_dispatch.h>
|
| 219 |
+
#include <ATen/ops/_unique2_compositeexplicitautograd_dispatch.h>
|
| 220 |
+
#include <ATen/ops/_unsafe_index_compositeexplicitautograd_dispatch.h>
|
| 221 |
+
#include <ATen/ops/_unsafe_index_put_compositeexplicitautograd_dispatch.h>
|
| 222 |
+
#include <ATen/ops/_unsafe_masked_index_compositeexplicitautograd_dispatch.h>
|
| 223 |
+
#include <ATen/ops/_unsafe_masked_index_put_accumulate_compositeexplicitautograd_dispatch.h>
|
| 224 |
+
#include <ATen/ops/_unsafe_view_compositeexplicitautograd_dispatch.h>
|
| 225 |
+
#include <ATen/ops/_values_copy_compositeexplicitautograd_dispatch.h>
|
| 226 |
+
#include <ATen/ops/_weight_norm_interface_compositeexplicitautograd_dispatch.h>
|
| 227 |
+
#include <ATen/ops/_weight_norm_interface_backward_compositeexplicitautograd_dispatch.h>
|
| 228 |
+
#include <ATen/ops/abs_compositeexplicitautograd_dispatch.h>
|
| 229 |
+
#include <ATen/ops/adaptive_avg_pool1d_compositeexplicitautograd_dispatch.h>
|
| 230 |
+
#include <ATen/ops/add_compositeexplicitautograd_dispatch.h>
|
| 231 |
+
#include <ATen/ops/addr_compositeexplicitautograd_dispatch.h>
|
| 232 |
+
#include <ATen/ops/affine_grid_generator_compositeexplicitautograd_dispatch.h>
|
| 233 |
+
#include <ATen/ops/alias_compositeexplicitautograd_dispatch.h>
|
| 234 |
+
#include <ATen/ops/alias_copy_compositeexplicitautograd_dispatch.h>
|
| 235 |
+
#include <ATen/ops/all_compositeexplicitautograd_dispatch.h>
|
| 236 |
+
#include <ATen/ops/allclose_compositeexplicitautograd_dispatch.h>
|
| 237 |
+
#include <ATen/ops/any_compositeexplicitautograd_dispatch.h>
|
| 238 |
+
#include <ATen/ops/arange_compositeexplicitautograd_dispatch.h>
|
| 239 |
+
#include <ATen/ops/as_strided_copy_compositeexplicitautograd_dispatch.h>
|
| 240 |
+
#include <ATen/ops/as_strided_scatter_compositeexplicitautograd_dispatch.h>
|
| 241 |
+
#include <ATen/ops/avg_pool1d_compositeexplicitautograd_dispatch.h>
|
| 242 |
+
#include <ATen/ops/bartlett_window_compositeexplicitautograd_dispatch.h>
|
| 243 |
+
#include <ATen/ops/batch_norm_backward_elemt_compositeexplicitautograd_dispatch.h>
|
| 244 |
+
#include <ATen/ops/batch_norm_backward_reduce_compositeexplicitautograd_dispatch.h>
|
| 245 |
+
#include <ATen/ops/batch_norm_gather_stats_compositeexplicitautograd_dispatch.h>
|
| 246 |
+
#include <ATen/ops/batch_norm_gather_stats_with_counts_compositeexplicitautograd_dispatch.h>
|
| 247 |
+
#include <ATen/ops/batch_norm_stats_compositeexplicitautograd_dispatch.h>
|
| 248 |
+
#include <ATen/ops/batch_norm_update_stats_compositeexplicitautograd_dispatch.h>
|
| 249 |
+
#include <ATen/ops/bernoulli_compositeexplicitautograd_dispatch.h>
|
| 250 |
+
#include <ATen/ops/binary_cross_entropy_with_logits_compositeexplicitautograd_dispatch.h>
|
| 251 |
+
#include <ATen/ops/bincount_compositeexplicitautograd_dispatch.h>
|
| 252 |
+
#include <ATen/ops/binomial_compositeexplicitautograd_dispatch.h>
|
| 253 |
+
#include <ATen/ops/bitwise_and_compositeexplicitautograd_dispatch.h>
|
| 254 |
+
#include <ATen/ops/bitwise_left_shift_compositeexplicitautograd_dispatch.h>
|
| 255 |
+
#include <ATen/ops/bitwise_or_compositeexplicitautograd_dispatch.h>
|
| 256 |
+
#include <ATen/ops/bitwise_right_shift_compositeexplicitautograd_dispatch.h>
|
| 257 |
+
#include <ATen/ops/bitwise_xor_compositeexplicitautograd_dispatch.h>
|
| 258 |
+
#include <ATen/ops/blackman_window_compositeexplicitautograd_dispatch.h>
|
| 259 |
+
#include <ATen/ops/block_diag_compositeexplicitautograd_dispatch.h>
|
| 260 |
+
#include <ATen/ops/bucketize_compositeexplicitautograd_dispatch.h>
|
| 261 |
+
#include <ATen/ops/cauchy_compositeexplicitautograd_dispatch.h>
|
| 262 |
+
#include <ATen/ops/ccol_indices_compositeexplicitautograd_dispatch.h>
|
| 263 |
+
#include <ATen/ops/ccol_indices_copy_compositeexplicitautograd_dispatch.h>
|
| 264 |
+
#include <ATen/ops/celu_compositeexplicitautograd_dispatch.h>
|
| 265 |
+
#include <ATen/ops/channel_shuffle_compositeexplicitautograd_dispatch.h>
|
| 266 |
+
#include <ATen/ops/cholesky_solve_compositeexplicitautograd_dispatch.h>
|
| 267 |
+
#include <ATen/ops/clone_compositeexplicitautograd_dispatch.h>
|
| 268 |
+
#include <ATen/ops/col_indices_compositeexplicitautograd_dispatch.h>
|
| 269 |
+
#include <ATen/ops/col_indices_copy_compositeexplicitautograd_dispatch.h>
|
| 270 |
+
#include <ATen/ops/complex_compositeexplicitautograd_dispatch.h>
|
| 271 |
+
#include <ATen/ops/conj_physical_compositeexplicitautograd_dispatch.h>
|
| 272 |
+
#include <ATen/ops/constant_pad_nd_compositeexplicitautograd_dispatch.h>
|
| 273 |
+
#include <ATen/ops/conv_depthwise3d_compositeexplicitautograd_dispatch.h>
|
| 274 |
+
#include <ATen/ops/conv_tbc_compositeexplicitautograd_dispatch.h>
|
| 275 |
+
#include <ATen/ops/convolution_compositeexplicitautograd_dispatch.h>
|
| 276 |
+
#include <ATen/ops/convolution_backward_compositeexplicitautograd_dispatch.h>
|
| 277 |
+
#include <ATen/ops/convolution_backward_overrideable_compositeexplicitautograd_dispatch.h>
|
| 278 |
+
#include <ATen/ops/convolution_overrideable_compositeexplicitautograd_dispatch.h>
|
| 279 |
+
#include <ATen/ops/copy_compositeexplicitautograd_dispatch.h>
|
| 280 |
+
#include <ATen/ops/copy_sparse_to_sparse_compositeexplicitautograd_dispatch.h>
|
| 281 |
+
#include <ATen/ops/copysign_compositeexplicitautograd_dispatch.h>
|
| 282 |
+
#include <ATen/ops/count_nonzero_compositeexplicitautograd_dispatch.h>
|
| 283 |
+
#include <ATen/ops/crow_indices_compositeexplicitautograd_dispatch.h>
|
| 284 |
+
#include <ATen/ops/crow_indices_copy_compositeexplicitautograd_dispatch.h>
|
| 285 |
+
#include <ATen/ops/cudnn_affine_grid_generator_compositeexplicitautograd_dispatch.h>
|
| 286 |
+
#include <ATen/ops/cudnn_affine_grid_generator_backward_compositeexplicitautograd_dispatch.h>
|
| 287 |
+
#include <ATen/ops/cudnn_batch_norm_compositeexplicitautograd_dispatch.h>
|
| 288 |
+
#include <ATen/ops/cudnn_batch_norm_backward_compositeexplicitautograd_dispatch.h>
|
| 289 |
+
#include <ATen/ops/cudnn_convolution_add_relu_compositeexplicitautograd_dispatch.h>
|
| 290 |
+
#include <ATen/ops/cudnn_convolution_relu_compositeexplicitautograd_dispatch.h>
|
| 291 |
+
#include <ATen/ops/cudnn_convolution_transpose_compositeexplicitautograd_dispatch.h>
|
| 292 |
+
#include <ATen/ops/cudnn_grid_sampler_compositeexplicitautograd_dispatch.h>
|
| 293 |
+
#include <ATen/ops/cudnn_grid_sampler_backward_compositeexplicitautograd_dispatch.h>
|
| 294 |
+
#include <ATen/ops/cummax_compositeexplicitautograd_dispatch.h>
|
| 295 |
+
#include <ATen/ops/cummin_compositeexplicitautograd_dispatch.h>
|
| 296 |
+
#include <ATen/ops/deg2rad_compositeexplicitautograd_dispatch.h>
|
| 297 |
+
#include <ATen/ops/dense_dim_compositeexplicitautograd_dispatch.h>
|
| 298 |
+
#include <ATen/ops/dequantize_compositeexplicitautograd_dispatch.h>
|
| 299 |
+
#include <ATen/ops/detach_compositeexplicitautograd_dispatch.h>
|
| 300 |
+
#include <ATen/ops/detach_copy_compositeexplicitautograd_dispatch.h>
|
| 301 |
+
#include <ATen/ops/diag_embed_compositeexplicitautograd_dispatch.h>
|
| 302 |
+
#include <ATen/ops/diagonal_compositeexplicitautograd_dispatch.h>
|
| 303 |
+
#include <ATen/ops/diagonal_backward_compositeexplicitautograd_dispatch.h>
|
| 304 |
+
#include <ATen/ops/diagonal_copy_compositeexplicitautograd_dispatch.h>
|
| 305 |
+
#include <ATen/ops/diagonal_scatter_compositeexplicitautograd_dispatch.h>
|
| 306 |
+
#include <ATen/ops/dist_compositeexplicitautograd_dispatch.h>
|
| 307 |
+
#include <ATen/ops/div_compositeexplicitautograd_dispatch.h>
|
| 308 |
+
#include <ATen/ops/dot_compositeexplicitautograd_dispatch.h>
|
| 309 |
+
#include <ATen/ops/embedding_compositeexplicitautograd_dispatch.h>
|
| 310 |
+
#include <ATen/ops/embedding_dense_backward_compositeexplicitautograd_dispatch.h>
|
| 311 |
+
#include <ATen/ops/embedding_renorm_compositeexplicitautograd_dispatch.h>
|
| 312 |
+
#include <ATen/ops/empty_compositeexplicitautograd_dispatch.h>
|
| 313 |
+
#include <ATen/ops/empty_like_compositeexplicitautograd_dispatch.h>
|
| 314 |
+
#include <ATen/ops/empty_permuted_compositeexplicitautograd_dispatch.h>
|
| 315 |
+
#include <ATen/ops/empty_quantized_compositeexplicitautograd_dispatch.h>
|
| 316 |
+
#include <ATen/ops/empty_strided_compositeexplicitautograd_dispatch.h>
|
| 317 |
+
#include <ATen/ops/expand_compositeexplicitautograd_dispatch.h>
|
| 318 |
+
#include <ATen/ops/expand_copy_compositeexplicitautograd_dispatch.h>
|
| 319 |
+
#include <ATen/ops/exponential_compositeexplicitautograd_dispatch.h>
|
| 320 |
+
#include <ATen/ops/eye_compositeexplicitautograd_dispatch.h>
|
| 321 |
+
#include <ATen/ops/fake_quantize_per_channel_affine_cachemask_compositeexplicitautograd_dispatch.h>
|
| 322 |
+
#include <ATen/ops/fake_quantize_per_tensor_affine_cachemask_compositeexplicitautograd_dispatch.h>
|
| 323 |
+
#include <ATen/ops/fft_fftfreq_compositeexplicitautograd_dispatch.h>
|
| 324 |
+
#include <ATen/ops/fft_rfftfreq_compositeexplicitautograd_dispatch.h>
|
| 325 |
+
#include <ATen/ops/fill_compositeexplicitautograd_dispatch.h>
|
| 326 |
+
#include <ATen/ops/flip_compositeexplicitautograd_dispatch.h>
|
| 327 |
+
#include <ATen/ops/floor_divide_compositeexplicitautograd_dispatch.h>
|
| 328 |
+
#include <ATen/ops/fmod_compositeexplicitautograd_dispatch.h>
|
| 329 |
+
#include <ATen/ops/frexp_compositeexplicitautograd_dispatch.h>
|
| 330 |
+
#include <ATen/ops/from_file_compositeexplicitautograd_dispatch.h>
|
| 331 |
+
#include <ATen/ops/full_compositeexplicitautograd_dispatch.h>
|
| 332 |
+
#include <ATen/ops/full_like_compositeexplicitautograd_dispatch.h>
|
| 333 |
+
#include <ATen/ops/geometric_compositeexplicitautograd_dispatch.h>
|
| 334 |
+
#include <ATen/ops/glu_backward_jvp_compositeexplicitautograd_dispatch.h>
|
| 335 |
+
#include <ATen/ops/glu_jvp_compositeexplicitautograd_dispatch.h>
|
| 336 |
+
#include <ATen/ops/grid_sampler_2d_compositeexplicitautograd_dispatch.h>
|
| 337 |
+
#include <ATen/ops/grid_sampler_2d_backward_compositeexplicitautograd_dispatch.h>
|
| 338 |
+
#include <ATen/ops/grid_sampler_3d_compositeexplicitautograd_dispatch.h>
|
| 339 |
+
#include <ATen/ops/grid_sampler_3d_backward_compositeexplicitautograd_dispatch.h>
|
| 340 |
+
#include <ATen/ops/hamming_window_compositeexplicitautograd_dispatch.h>
|
| 341 |
+
#include <ATen/ops/hann_window_compositeexplicitautograd_dispatch.h>
|
| 342 |
+
#include <ATen/ops/hardswish_backward_compositeexplicitautograd_dispatch.h>
|
| 343 |
+
#include <ATen/ops/huber_loss_backward_compositeexplicitautograd_dispatch.h>
|
| 344 |
+
#include <ATen/ops/index_fill_compositeexplicitautograd_dispatch.h>
|
| 345 |
+
#include <ATen/ops/index_put_compositeexplicitautograd_dispatch.h>
|
| 346 |
+
#include <ATen/ops/indices_compositeexplicitautograd_dispatch.h>
|
| 347 |
+
#include <ATen/ops/indices_copy_compositeexplicitautograd_dispatch.h>
|
| 348 |
+
#include <ATen/ops/int_repr_compositeexplicitautograd_dispatch.h>
|
| 349 |
+
#include <ATen/ops/is_coalesced_compositeexplicitautograd_dispatch.h>
|
| 350 |
+
#include <ATen/ops/is_pinned_compositeexplicitautograd_dispatch.h>
|
| 351 |
+
#include <ATen/ops/is_same_size_compositeexplicitautograd_dispatch.h>
|
| 352 |
+
#include <ATen/ops/isinf_compositeexplicitautograd_dispatch.h>
|
| 353 |
+
#include <ATen/ops/isnan_compositeexplicitautograd_dispatch.h>
|
| 354 |
+
#include <ATen/ops/kaiser_window_compositeexplicitautograd_dispatch.h>
|
| 355 |
+
#include <ATen/ops/kthvalue_compositeexplicitautograd_dispatch.h>
|
| 356 |
+
#include <ATen/ops/lift_compositeexplicitautograd_dispatch.h>
|
| 357 |
+
#include <ATen/ops/lift_fresh_compositeexplicitautograd_dispatch.h>
|
| 358 |
+
#include <ATen/ops/lift_fresh_copy_compositeexplicitautograd_dispatch.h>
|
| 359 |
+
#include <ATen/ops/linalg_lstsq_compositeexplicitautograd_dispatch.h>
|
| 360 |
+
#include <ATen/ops/linalg_matrix_exp_compositeexplicitautograd_dispatch.h>
|
| 361 |
+
#include <ATen/ops/linalg_pinv_compositeexplicitautograd_dispatch.h>
|
| 362 |
+
#include <ATen/ops/linear_compositeexplicitautograd_dispatch.h>
|
| 363 |
+
#include <ATen/ops/linear_backward_compositeexplicitautograd_dispatch.h>
|
| 364 |
+
#include <ATen/ops/linspace_compositeexplicitautograd_dispatch.h>
|
| 365 |
+
#include <ATen/ops/log_normal_compositeexplicitautograd_dispatch.h>
|
| 366 |
+
#include <ATen/ops/log_softmax_compositeexplicitautograd_dispatch.h>
|
| 367 |
+
#include <ATen/ops/logcumsumexp_compositeexplicitautograd_dispatch.h>
|
| 368 |
+
#include <ATen/ops/logical_and_compositeexplicitautograd_dispatch.h>
|
| 369 |
+
#include <ATen/ops/logical_not_compositeexplicitautograd_dispatch.h>
|
| 370 |
+
#include <ATen/ops/logical_or_compositeexplicitautograd_dispatch.h>
|
| 371 |
+
#include <ATen/ops/logical_xor_compositeexplicitautograd_dispatch.h>
|
| 372 |
+
#include <ATen/ops/logspace_compositeexplicitautograd_dispatch.h>
|
| 373 |
+
#include <ATen/ops/logsumexp_compositeexplicitautograd_dispatch.h>
|
| 374 |
+
#include <ATen/ops/lshift_compositeexplicitautograd_dispatch.h>
|
| 375 |
+
#include <ATen/ops/lstm_mps_backward_compositeexplicitautograd_dispatch.h>
|
| 376 |
+
#include <ATen/ops/masked_fill_compositeexplicitautograd_dispatch.h>
|
| 377 |
+
#include <ATen/ops/masked_scatter_compositeexplicitautograd_dispatch.h>
|
| 378 |
+
#include <ATen/ops/masked_scatter_backward_compositeexplicitautograd_dispatch.h>
|
| 379 |
+
#include <ATen/ops/matmul_backward_compositeexplicitautograd_dispatch.h>
|
| 380 |
+
#include <ATen/ops/max_pool2d_backward_compositeexplicitautograd_dispatch.h>
|
| 381 |
+
#include <ATen/ops/mean_compositeexplicitautograd_dispatch.h>
|
| 382 |
+
#include <ATen/ops/median_compositeexplicitautograd_dispatch.h>
|
| 383 |
+
#include <ATen/ops/miopen_batch_norm_compositeexplicitautograd_dispatch.h>
|
| 384 |
+
#include <ATen/ops/miopen_batch_norm_backward_compositeexplicitautograd_dispatch.h>
|
| 385 |
+
#include <ATen/ops/miopen_convolution_compositeexplicitautograd_dispatch.h>
|
| 386 |
+
#include <ATen/ops/miopen_convolution_transpose_compositeexplicitautograd_dispatch.h>
|
| 387 |
+
#include <ATen/ops/miopen_depthwise_convolution_compositeexplicitautograd_dispatch.h>
|
| 388 |
+
#include <ATen/ops/miopen_rnn_compositeexplicitautograd_dispatch.h>
|
| 389 |
+
#include <ATen/ops/miopen_rnn_backward_compositeexplicitautograd_dispatch.h>
|
| 390 |
+
#include <ATen/ops/mkldnn_adaptive_avg_pool2d_backward_compositeexplicitautograd_dispatch.h>
|
| 391 |
+
#include <ATen/ops/mkldnn_convolution_compositeexplicitautograd_dispatch.h>
|
| 392 |
+
#include <ATen/ops/mkldnn_linear_compositeexplicitautograd_dispatch.h>
|
| 393 |
+
#include <ATen/ops/mkldnn_linear_backward_compositeexplicitautograd_dispatch.h>
|
| 394 |
+
#include <ATen/ops/mkldnn_linear_backward_input_compositeexplicitautograd_dispatch.h>
|
| 395 |
+
#include <ATen/ops/mkldnn_linear_backward_weights_compositeexplicitautograd_dispatch.h>
|
| 396 |
+
#include <ATen/ops/mkldnn_max_pool2d_compositeexplicitautograd_dispatch.h>
|
| 397 |
+
#include <ATen/ops/mkldnn_max_pool2d_backward_compositeexplicitautograd_dispatch.h>
|
| 398 |
+
#include <ATen/ops/mkldnn_max_pool3d_compositeexplicitautograd_dispatch.h>
|
| 399 |
+
#include <ATen/ops/mkldnn_max_pool3d_backward_compositeexplicitautograd_dispatch.h>
|
| 400 |
+
#include <ATen/ops/mkldnn_reorder_conv2d_weight_compositeexplicitautograd_dispatch.h>
|
| 401 |
+
#include <ATen/ops/mkldnn_reorder_conv3d_weight_compositeexplicitautograd_dispatch.h>
|
| 402 |
+
#include <ATen/ops/mkldnn_rnn_layer_compositeexplicitautograd_dispatch.h>
|
| 403 |
+
#include <ATen/ops/mkldnn_rnn_layer_backward_compositeexplicitautograd_dispatch.h>
|
| 404 |
+
#include <ATen/ops/mode_compositeexplicitautograd_dispatch.h>
|
| 405 |
+
#include <ATen/ops/mps_convolution_backward_compositeexplicitautograd_dispatch.h>
|
| 406 |
+
#include <ATen/ops/mps_convolution_transpose_backward_compositeexplicitautograd_dispatch.h>
|
| 407 |
+
#include <ATen/ops/mul_compositeexplicitautograd_dispatch.h>
|
| 408 |
+
#include <ATen/ops/mv_compositeexplicitautograd_dispatch.h>
|
| 409 |
+
#include <ATen/ops/mvlgamma_compositeexplicitautograd_dispatch.h>
|
| 410 |
+
#include <ATen/ops/nan_to_num_compositeexplicitautograd_dispatch.h>
|
| 411 |
+
#include <ATen/ops/nanmedian_compositeexplicitautograd_dispatch.h>
|
| 412 |
+
#include <ATen/ops/native_batch_norm_backward_compositeexplicitautograd_dispatch.h>
|
| 413 |
+
#include <ATen/ops/native_dropout_compositeexplicitautograd_dispatch.h>
|
| 414 |
+
#include <ATen/ops/native_dropout_backward_compositeexplicitautograd_dispatch.h>
|
| 415 |
+
#include <ATen/ops/native_group_norm_compositeexplicitautograd_dispatch.h>
|
| 416 |
+
#include <ATen/ops/native_group_norm_backward_compositeexplicitautograd_dispatch.h>
|
| 417 |
+
#include <ATen/ops/native_layer_norm_compositeexplicitautograd_dispatch.h>
|
| 418 |
+
#include <ATen/ops/native_layer_norm_backward_compositeexplicitautograd_dispatch.h>
|
| 419 |
+
#include <ATen/ops/native_norm_compositeexplicitautograd_dispatch.h>
|
| 420 |
+
#include <ATen/ops/new_empty_compositeexplicitautograd_dispatch.h>
|
| 421 |
+
#include <ATen/ops/new_empty_strided_compositeexplicitautograd_dispatch.h>
|
| 422 |
+
#include <ATen/ops/new_full_compositeexplicitautograd_dispatch.h>
|
| 423 |
+
#include <ATen/ops/new_ones_compositeexplicitautograd_dispatch.h>
|
| 424 |
+
#include <ATen/ops/new_zeros_compositeexplicitautograd_dispatch.h>
|
| 425 |
+
#include <ATen/ops/norm_compositeexplicitautograd_dispatch.h>
|
| 426 |
+
#include <ATen/ops/normal_compositeexplicitautograd_dispatch.h>
|
| 427 |
+
#include <ATen/ops/ones_compositeexplicitautograd_dispatch.h>
|
| 428 |
+
#include <ATen/ops/ones_like_compositeexplicitautograd_dispatch.h>
|
| 429 |
+
#include <ATen/ops/permute_compositeexplicitautograd_dispatch.h>
|
| 430 |
+
#include <ATen/ops/permute_copy_compositeexplicitautograd_dispatch.h>
|
| 431 |
+
#include <ATen/ops/pixel_shuffle_compositeexplicitautograd_dispatch.h>
|
| 432 |
+
#include <ATen/ops/pixel_unshuffle_compositeexplicitautograd_dispatch.h>
|
| 433 |
+
#include <ATen/ops/poisson_compositeexplicitautograd_dispatch.h>
|
| 434 |
+
#include <ATen/ops/polar_compositeexplicitautograd_dispatch.h>
|
| 435 |
+
#include <ATen/ops/polygamma_compositeexplicitautograd_dispatch.h>
|
| 436 |
+
#include <ATen/ops/prod_compositeexplicitautograd_dispatch.h>
|
| 437 |
+
#include <ATen/ops/put_compositeexplicitautograd_dispatch.h>
|
| 438 |
+
#include <ATen/ops/q_per_channel_scales_compositeexplicitautograd_dispatch.h>
|
| 439 |
+
#include <ATen/ops/q_per_channel_zero_points_compositeexplicitautograd_dispatch.h>
|
| 440 |
+
#include <ATen/ops/quantize_per_channel_compositeexplicitautograd_dispatch.h>
|
| 441 |
+
#include <ATen/ops/quantize_per_tensor_compositeexplicitautograd_dispatch.h>
|
| 442 |
+
#include <ATen/ops/quantize_per_tensor_dynamic_compositeexplicitautograd_dispatch.h>
|
| 443 |
+
#include <ATen/ops/quantized_batch_norm_compositeexplicitautograd_dispatch.h>
|
| 444 |
+
#include <ATen/ops/quantized_max_pool1d_compositeexplicitautograd_dispatch.h>
|
| 445 |
+
#include <ATen/ops/quantized_max_pool2d_compositeexplicitautograd_dispatch.h>
|
| 446 |
+
#include <ATen/ops/quantized_max_pool3d_compositeexplicitautograd_dispatch.h>
|
| 447 |
+
#include <ATen/ops/rad2deg_compositeexplicitautograd_dispatch.h>
|
| 448 |
+
#include <ATen/ops/rand_compositeexplicitautograd_dispatch.h>
|
| 449 |
+
#include <ATen/ops/rand_like_compositeexplicitautograd_dispatch.h>
|
| 450 |
+
#include <ATen/ops/randint_compositeexplicitautograd_dispatch.h>
|
| 451 |
+
#include <ATen/ops/randint_like_compositeexplicitautograd_dispatch.h>
|
| 452 |
+
#include <ATen/ops/randn_compositeexplicitautograd_dispatch.h>
|
| 453 |
+
#include <ATen/ops/randn_like_compositeexplicitautograd_dispatch.h>
|
| 454 |
+
#include <ATen/ops/random_compositeexplicitautograd_dispatch.h>
|
| 455 |
+
#include <ATen/ops/randperm_compositeexplicitautograd_dispatch.h>
|
| 456 |
+
#include <ATen/ops/range_compositeexplicitautograd_dispatch.h>
|
| 457 |
+
#include <ATen/ops/relu_compositeexplicitautograd_dispatch.h>
|
| 458 |
+
#include <ATen/ops/remainder_compositeexplicitautograd_dispatch.h>
|
| 459 |
+
#include <ATen/ops/repeat_compositeexplicitautograd_dispatch.h>
|
| 460 |
+
#include <ATen/ops/repeat_interleave_compositeexplicitautograd_dispatch.h>
|
| 461 |
+
#include <ATen/ops/resize_compositeexplicitautograd_dispatch.h>
|
| 462 |
+
#include <ATen/ops/resize_as_compositeexplicitautograd_dispatch.h>
|
| 463 |
+
#include <ATen/ops/resize_as_sparse_compositeexplicitautograd_dispatch.h>
|
| 464 |
+
#include <ATen/ops/roll_compositeexplicitautograd_dispatch.h>
|
| 465 |
+
#include <ATen/ops/rot90_compositeexplicitautograd_dispatch.h>
|
| 466 |
+
#include <ATen/ops/row_indices_compositeexplicitautograd_dispatch.h>
|
| 467 |
+
#include <ATen/ops/row_indices_copy_compositeexplicitautograd_dispatch.h>
|
| 468 |
+
#include <ATen/ops/rrelu_with_noise_compositeexplicitautograd_dispatch.h>
|
| 469 |
+
#include <ATen/ops/rrelu_with_noise_backward_compositeexplicitautograd_dispatch.h>
|
| 470 |
+
#include <ATen/ops/rshift_compositeexplicitautograd_dispatch.h>
|
| 471 |
+
#include <ATen/ops/rsub_compositeexplicitautograd_dispatch.h>
|
| 472 |
+
#include <ATen/ops/scalar_tensor_compositeexplicitautograd_dispatch.h>
|
| 473 |
+
#include <ATen/ops/segment_reduce_compositeexplicitautograd_dispatch.h>
|
| 474 |
+
#include <ATen/ops/select_compositeexplicitautograd_dispatch.h>
|
| 475 |
+
#include <ATen/ops/select_backward_compositeexplicitautograd_dispatch.h>
|
| 476 |
+
#include <ATen/ops/select_copy_compositeexplicitautograd_dispatch.h>
|
| 477 |
+
#include <ATen/ops/select_scatter_compositeexplicitautograd_dispatch.h>
|
| 478 |
+
#include <ATen/ops/set_compositeexplicitautograd_dispatch.h>
|
| 479 |
+
#include <ATen/ops/slice_compositeexplicitautograd_dispatch.h>
|
| 480 |
+
#include <ATen/ops/slice_backward_compositeexplicitautograd_dispatch.h>
|
| 481 |
+
#include <ATen/ops/slice_copy_compositeexplicitautograd_dispatch.h>
|
| 482 |
+
#include <ATen/ops/slice_inverse_compositeexplicitautograd_dispatch.h>
|
| 483 |
+
#include <ATen/ops/slice_scatter_compositeexplicitautograd_dispatch.h>
|
| 484 |
+
#include <ATen/ops/slow_conv_dilated2d_compositeexplicitautograd_dispatch.h>
|
| 485 |
+
#include <ATen/ops/slow_conv_dilated3d_compositeexplicitautograd_dispatch.h>
|
| 486 |
+
#include <ATen/ops/smooth_l1_loss_backward_compositeexplicitautograd_dispatch.h>
|
| 487 |
+
#include <ATen/ops/soft_margin_loss_compositeexplicitautograd_dispatch.h>
|
| 488 |
+
#include <ATen/ops/soft_margin_loss_backward_compositeexplicitautograd_dispatch.h>
|
| 489 |
+
#include <ATen/ops/softmax_compositeexplicitautograd_dispatch.h>
|
| 490 |
+
#include <ATen/ops/sort_compositeexplicitautograd_dispatch.h>
|
| 491 |
+
#include <ATen/ops/sparse_compressed_tensor_compositeexplicitautograd_dispatch.h>
|
| 492 |
+
#include <ATen/ops/sparse_coo_tensor_compositeexplicitautograd_dispatch.h>
|
| 493 |
+
#include <ATen/ops/sparse_dim_compositeexplicitautograd_dispatch.h>
|
| 494 |
+
#include <ATen/ops/sparse_mask_compositeexplicitautograd_dispatch.h>
|
| 495 |
+
#include <ATen/ops/sparse_resize_compositeexplicitautograd_dispatch.h>
|
| 496 |
+
#include <ATen/ops/sparse_resize_and_clear_compositeexplicitautograd_dispatch.h>
|
| 497 |
+
#include <ATen/ops/special_chebyshev_polynomial_t_compositeexplicitautograd_dispatch.h>
|
| 498 |
+
#include <ATen/ops/special_chebyshev_polynomial_u_compositeexplicitautograd_dispatch.h>
|
| 499 |
+
#include <ATen/ops/special_chebyshev_polynomial_v_compositeexplicitautograd_dispatch.h>
|
| 500 |
+
#include <ATen/ops/special_chebyshev_polynomial_w_compositeexplicitautograd_dispatch.h>
|
| 501 |
+
#include <ATen/ops/special_hermite_polynomial_h_compositeexplicitautograd_dispatch.h>
|
| 502 |
+
#include <ATen/ops/special_hermite_polynomial_he_compositeexplicitautograd_dispatch.h>
|
| 503 |
+
#include <ATen/ops/special_laguerre_polynomial_l_compositeexplicitautograd_dispatch.h>
|
| 504 |
+
#include <ATen/ops/special_legendre_polynomial_p_compositeexplicitautograd_dispatch.h>
|
| 505 |
+
#include <ATen/ops/special_shifted_chebyshev_polynomial_t_compositeexplicitautograd_dispatch.h>
|
| 506 |
+
#include <ATen/ops/special_shifted_chebyshev_polynomial_u_compositeexplicitautograd_dispatch.h>
|
| 507 |
+
#include <ATen/ops/special_shifted_chebyshev_polynomial_v_compositeexplicitautograd_dispatch.h>
|
| 508 |
+
#include <ATen/ops/special_shifted_chebyshev_polynomial_w_compositeexplicitautograd_dispatch.h>
|
| 509 |
+
#include <ATen/ops/special_xlog1py_compositeexplicitautograd_dispatch.h>
|
| 510 |
+
#include <ATen/ops/special_zeta_compositeexplicitautograd_dispatch.h>
|
| 511 |
+
#include <ATen/ops/split_compositeexplicitautograd_dispatch.h>
|
| 512 |
+
#include <ATen/ops/split_copy_compositeexplicitautograd_dispatch.h>
|
| 513 |
+
#include <ATen/ops/split_with_sizes_compositeexplicitautograd_dispatch.h>
|
| 514 |
+
#include <ATen/ops/split_with_sizes_copy_compositeexplicitautograd_dispatch.h>
|
| 515 |
+
#include <ATen/ops/squeeze_compositeexplicitautograd_dispatch.h>
|
| 516 |
+
#include <ATen/ops/squeeze_copy_compositeexplicitautograd_dispatch.h>
|
| 517 |
+
#include <ATen/ops/stack_compositeexplicitautograd_dispatch.h>
|
| 518 |
+
#include <ATen/ops/std_mean_compositeexplicitautograd_dispatch.h>
|
| 519 |
+
#include <ATen/ops/sub_compositeexplicitautograd_dispatch.h>
|
| 520 |
+
#include <ATen/ops/sum_compositeexplicitautograd_dispatch.h>
|
| 521 |
+
#include <ATen/ops/sym_constrain_range_compositeexplicitautograd_dispatch.h>
|
| 522 |
+
#include <ATen/ops/sym_constrain_range_for_size_compositeexplicitautograd_dispatch.h>
|
| 523 |
+
#include <ATen/ops/t_compositeexplicitautograd_dispatch.h>
|
| 524 |
+
#include <ATen/ops/t_copy_compositeexplicitautograd_dispatch.h>
|
| 525 |
+
#include <ATen/ops/to_mkldnn_compositeexplicitautograd_dispatch.h>
|
| 526 |
+
#include <ATen/ops/to_padded_tensor_compositeexplicitautograd_dispatch.h>
|
| 527 |
+
#include <ATen/ops/trace_compositeexplicitautograd_dispatch.h>
|
| 528 |
+
#include <ATen/ops/transpose_compositeexplicitautograd_dispatch.h>
|
| 529 |
+
#include <ATen/ops/transpose_copy_compositeexplicitautograd_dispatch.h>
|
| 530 |
+
#include <ATen/ops/tril_indices_compositeexplicitautograd_dispatch.h>
|
| 531 |
+
#include <ATen/ops/triu_indices_compositeexplicitautograd_dispatch.h>
|
| 532 |
+
#include <ATen/ops/unbind_compositeexplicitautograd_dispatch.h>
|
| 533 |
+
#include <ATen/ops/unbind_copy_compositeexplicitautograd_dispatch.h>
|
| 534 |
+
#include <ATen/ops/unfold_backward_compositeexplicitautograd_dispatch.h>
|
| 535 |
+
#include <ATen/ops/unfold_copy_compositeexplicitautograd_dispatch.h>
|
| 536 |
+
#include <ATen/ops/uniform_compositeexplicitautograd_dispatch.h>
|
| 537 |
+
#include <ATen/ops/unique_consecutive_compositeexplicitautograd_dispatch.h>
|
| 538 |
+
#include <ATen/ops/unique_dim_compositeexplicitautograd_dispatch.h>
|
| 539 |
+
#include <ATen/ops/unique_dim_consecutive_compositeexplicitautograd_dispatch.h>
|
| 540 |
+
#include <ATen/ops/unsafe_split_compositeexplicitautograd_dispatch.h>
|
| 541 |
+
#include <ATen/ops/unsafe_split_with_sizes_compositeexplicitautograd_dispatch.h>
|
| 542 |
+
#include <ATen/ops/unsqueeze_compositeexplicitautograd_dispatch.h>
|
| 543 |
+
#include <ATen/ops/unsqueeze_copy_compositeexplicitautograd_dispatch.h>
|
| 544 |
+
#include <ATen/ops/upsample_bilinear2d_compositeexplicitautograd_dispatch.h>
|
| 545 |
+
#include <ATen/ops/upsample_nearest2d_compositeexplicitautograd_dispatch.h>
|
| 546 |
+
#include <ATen/ops/values_compositeexplicitautograd_dispatch.h>
|
| 547 |
+
#include <ATen/ops/values_copy_compositeexplicitautograd_dispatch.h>
|
| 548 |
+
#include <ATen/ops/var_mean_compositeexplicitautograd_dispatch.h>
|
| 549 |
+
#include <ATen/ops/vdot_compositeexplicitautograd_dispatch.h>
|
| 550 |
+
#include <ATen/ops/view_compositeexplicitautograd_dispatch.h>
|
| 551 |
+
#include <ATen/ops/view_as_complex_copy_compositeexplicitautograd_dispatch.h>
|
| 552 |
+
#include <ATen/ops/view_as_real_copy_compositeexplicitautograd_dispatch.h>
|
| 553 |
+
#include <ATen/ops/view_copy_compositeexplicitautograd_dispatch.h>
|
| 554 |
+
#include <ATen/ops/xlogy_compositeexplicitautograd_dispatch.h>
|
| 555 |
+
#include <ATen/ops/zero_compositeexplicitautograd_dispatch.h>
|
| 556 |
+
#include <ATen/ops/zeros_compositeexplicitautograd_dispatch.h>
|
| 557 |
+
#include <ATen/ops/zeros_like_compositeexplicitautograd_dispatch.h>
|
| 558 |
+
|
| 559 |
+
|
| 560 |
+
|
lib/python3.12/site-packages/torch/include/ATen/CompositeImplicitAutogradFunctions_inl.h
ADDED
|
@@ -0,0 +1,501 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
// @generated by torchgen/gen.py from DispatchKeyFunctions_inl.h
|
| 3 |
+
|
| 4 |
+
// NB: The implementing C++ file is RegisterDispatchKey.cpp
|
| 5 |
+
|
| 6 |
+
// The only #includes we need are for custom classes that have defaults in the C++ API
|
| 7 |
+
#include <c10/core/MemoryFormat.h>
|
| 8 |
+
#include <c10/core/Scalar.h>
|
| 9 |
+
#include <ATen/core/Reduction.h>
|
| 10 |
+
|
| 11 |
+
#if defined(AT_PER_OPERATOR_HEADERS) && defined(TORCH_ASSERT_ONLY_METHOD_OPERATORS)
|
| 12 |
+
#error This change adds a dependency on all pytorch operators, meaning the \
|
| 13 |
+
file will need to be re-compiled every time an operator is changed or added. \
|
| 14 |
+
Consider including a specific operator from \
|
| 15 |
+
<ATen/ops/{my_operator}_compositeimplicitautograd_dispatch.h>. \
|
| 16 |
+
See NOTE [TORCH_ASSERT_ONLY_METHOD_OPERATORS].
|
| 17 |
+
#endif
|
| 18 |
+
|
| 19 |
+
#include <ATen/ops/_add_batch_dim_compositeimplicitautograd_dispatch.h>
|
| 20 |
+
#include <ATen/ops/_autocast_to_full_precision_compositeimplicitautograd_dispatch.h>
|
| 21 |
+
#include <ATen/ops/_autocast_to_reduced_precision_compositeimplicitautograd_dispatch.h>
|
| 22 |
+
#include <ATen/ops/_backward_compositeimplicitautograd_dispatch.h>
|
| 23 |
+
#include <ATen/ops/_batch_norm_impl_index_compositeimplicitautograd_dispatch.h>
|
| 24 |
+
#include <ATen/ops/_batch_norm_impl_index_backward_compositeimplicitautograd_dispatch.h>
|
| 25 |
+
#include <ATen/ops/_cast_Byte_compositeimplicitautograd_dispatch.h>
|
| 26 |
+
#include <ATen/ops/_cast_Char_compositeimplicitautograd_dispatch.h>
|
| 27 |
+
#include <ATen/ops/_cast_Double_compositeimplicitautograd_dispatch.h>
|
| 28 |
+
#include <ATen/ops/_cast_Float_compositeimplicitautograd_dispatch.h>
|
| 29 |
+
#include <ATen/ops/_cast_Half_compositeimplicitautograd_dispatch.h>
|
| 30 |
+
#include <ATen/ops/_cast_Int_compositeimplicitautograd_dispatch.h>
|
| 31 |
+
#include <ATen/ops/_cast_Long_compositeimplicitautograd_dispatch.h>
|
| 32 |
+
#include <ATen/ops/_cast_Short_compositeimplicitautograd_dispatch.h>
|
| 33 |
+
#include <ATen/ops/_choose_qparams_per_tensor_compositeimplicitautograd_dispatch.h>
|
| 34 |
+
#include <ATen/ops/_convolution_compositeimplicitautograd_dispatch.h>
|
| 35 |
+
#include <ATen/ops/_convolution_double_backward_compositeimplicitautograd_dispatch.h>
|
| 36 |
+
#include <ATen/ops/_convolution_mode_compositeimplicitautograd_dispatch.h>
|
| 37 |
+
#include <ATen/ops/_cufft_clear_plan_cache_compositeimplicitautograd_dispatch.h>
|
| 38 |
+
#include <ATen/ops/_cufft_get_plan_cache_max_size_compositeimplicitautograd_dispatch.h>
|
| 39 |
+
#include <ATen/ops/_cufft_get_plan_cache_size_compositeimplicitautograd_dispatch.h>
|
| 40 |
+
#include <ATen/ops/_cufft_set_plan_cache_max_size_compositeimplicitautograd_dispatch.h>
|
| 41 |
+
#include <ATen/ops/_debug_has_internal_overlap_compositeimplicitautograd_dispatch.h>
|
| 42 |
+
#include <ATen/ops/_dim_arange_compositeimplicitautograd_dispatch.h>
|
| 43 |
+
#include <ATen/ops/_embedding_bag_sparse_backward_compositeimplicitautograd_dispatch.h>
|
| 44 |
+
#include <ATen/ops/_gather_sparse_backward_compositeimplicitautograd_dispatch.h>
|
| 45 |
+
#include <ATen/ops/_grid_sampler_2d_cpu_fallback_backward_compositeimplicitautograd_dispatch.h>
|
| 46 |
+
#include <ATen/ops/_has_compatible_shallow_copy_type_compositeimplicitautograd_dispatch.h>
|
| 47 |
+
#include <ATen/ops/_is_zerotensor_compositeimplicitautograd_dispatch.h>
|
| 48 |
+
#include <ATen/ops/_lu_with_info_compositeimplicitautograd_dispatch.h>
|
| 49 |
+
#include <ATen/ops/_nnpack_available_compositeimplicitautograd_dispatch.h>
|
| 50 |
+
#include <ATen/ops/_pack_padded_sequence_backward_compositeimplicitautograd_dispatch.h>
|
| 51 |
+
#include <ATen/ops/_pad_circular_compositeimplicitautograd_dispatch.h>
|
| 52 |
+
#include <ATen/ops/_pad_enum_compositeimplicitautograd_dispatch.h>
|
| 53 |
+
#include <ATen/ops/_pad_packed_sequence_compositeimplicitautograd_dispatch.h>
|
| 54 |
+
#include <ATen/ops/_propagate_xla_data_compositeimplicitautograd_dispatch.h>
|
| 55 |
+
#include <ATen/ops/_remove_batch_dim_compositeimplicitautograd_dispatch.h>
|
| 56 |
+
#include <ATen/ops/_reshape_from_tensor_compositeimplicitautograd_dispatch.h>
|
| 57 |
+
#include <ATen/ops/_rowwise_prune_compositeimplicitautograd_dispatch.h>
|
| 58 |
+
#include <ATen/ops/_saturate_weight_to_fp16_compositeimplicitautograd_dispatch.h>
|
| 59 |
+
#include <ATen/ops/_scaled_dot_product_attention_math_compositeimplicitautograd_dispatch.h>
|
| 60 |
+
#include <ATen/ops/_shape_as_tensor_compositeimplicitautograd_dispatch.h>
|
| 61 |
+
#include <ATen/ops/_sobol_engine_draw_compositeimplicitautograd_dispatch.h>
|
| 62 |
+
#include <ATen/ops/_sobol_engine_ff_compositeimplicitautograd_dispatch.h>
|
| 63 |
+
#include <ATen/ops/_sobol_engine_initialize_state_compositeimplicitautograd_dispatch.h>
|
| 64 |
+
#include <ATen/ops/_sobol_engine_scramble_compositeimplicitautograd_dispatch.h>
|
| 65 |
+
#include <ATen/ops/_sparse_bsc_tensor_unsafe_compositeimplicitautograd_dispatch.h>
|
| 66 |
+
#include <ATen/ops/_sparse_bsr_tensor_unsafe_compositeimplicitautograd_dispatch.h>
|
| 67 |
+
#include <ATen/ops/_sparse_compressed_tensor_unsafe_compositeimplicitautograd_dispatch.h>
|
| 68 |
+
#include <ATen/ops/_sparse_coo_tensor_unsafe_compositeimplicitautograd_dispatch.h>
|
| 69 |
+
#include <ATen/ops/_sparse_csc_tensor_unsafe_compositeimplicitautograd_dispatch.h>
|
| 70 |
+
#include <ATen/ops/_sparse_csr_tensor_unsafe_compositeimplicitautograd_dispatch.h>
|
| 71 |
+
#include <ATen/ops/_sparse_log_softmax_compositeimplicitautograd_dispatch.h>
|
| 72 |
+
#include <ATen/ops/_sparse_mm_compositeimplicitautograd_dispatch.h>
|
| 73 |
+
#include <ATen/ops/_sparse_softmax_compositeimplicitautograd_dispatch.h>
|
| 74 |
+
#include <ATen/ops/_sparse_sum_compositeimplicitautograd_dispatch.h>
|
| 75 |
+
#include <ATen/ops/_test_ambiguous_defaults_compositeimplicitautograd_dispatch.h>
|
| 76 |
+
#include <ATen/ops/_test_autograd_multiple_dispatch_compositeimplicitautograd_dispatch.h>
|
| 77 |
+
#include <ATen/ops/_test_check_tensor_compositeimplicitautograd_dispatch.h>
|
| 78 |
+
#include <ATen/ops/_test_serialization_subcmul_compositeimplicitautograd_dispatch.h>
|
| 79 |
+
#include <ATen/ops/_test_string_default_compositeimplicitautograd_dispatch.h>
|
| 80 |
+
#include <ATen/ops/_thnn_differentiable_gru_cell_backward_compositeimplicitautograd_dispatch.h>
|
| 81 |
+
#include <ATen/ops/_thnn_differentiable_lstm_cell_backward_compositeimplicitautograd_dispatch.h>
|
| 82 |
+
#include <ATen/ops/_thnn_fused_lstm_cell_backward_compositeimplicitautograd_dispatch.h>
|
| 83 |
+
#include <ATen/ops/_to_cpu_compositeimplicitautograd_dispatch.h>
|
| 84 |
+
#include <ATen/ops/_unpack_dual_compositeimplicitautograd_dispatch.h>
|
| 85 |
+
#include <ATen/ops/_upsample_bicubic2d_aa_compositeimplicitautograd_dispatch.h>
|
| 86 |
+
#include <ATen/ops/_upsample_bilinear2d_aa_compositeimplicitautograd_dispatch.h>
|
| 87 |
+
#include <ATen/ops/_upsample_nearest_exact1d_compositeimplicitautograd_dispatch.h>
|
| 88 |
+
#include <ATen/ops/_upsample_nearest_exact2d_compositeimplicitautograd_dispatch.h>
|
| 89 |
+
#include <ATen/ops/_upsample_nearest_exact3d_compositeimplicitautograd_dispatch.h>
|
| 90 |
+
#include <ATen/ops/_use_cudnn_rnn_flatten_weight_compositeimplicitautograd_dispatch.h>
|
| 91 |
+
#include <ATen/ops/_validate_sparse_bsc_tensor_args_compositeimplicitautograd_dispatch.h>
|
| 92 |
+
#include <ATen/ops/_validate_sparse_bsr_tensor_args_compositeimplicitautograd_dispatch.h>
|
| 93 |
+
#include <ATen/ops/_validate_sparse_compressed_tensor_args_compositeimplicitautograd_dispatch.h>
|
| 94 |
+
#include <ATen/ops/_validate_sparse_coo_tensor_args_compositeimplicitautograd_dispatch.h>
|
| 95 |
+
#include <ATen/ops/_validate_sparse_csc_tensor_args_compositeimplicitautograd_dispatch.h>
|
| 96 |
+
#include <ATen/ops/_validate_sparse_csr_tensor_args_compositeimplicitautograd_dispatch.h>
|
| 97 |
+
#include <ATen/ops/_version_compositeimplicitautograd_dispatch.h>
|
| 98 |
+
#include <ATen/ops/_weight_norm_compositeimplicitautograd_dispatch.h>
|
| 99 |
+
#include <ATen/ops/_weight_norm_differentiable_backward_compositeimplicitautograd_dispatch.h>
|
| 100 |
+
#include <ATen/ops/_wrapped_linear_prepack_compositeimplicitautograd_dispatch.h>
|
| 101 |
+
#include <ATen/ops/_wrapped_quantized_linear_prepacked_compositeimplicitautograd_dispatch.h>
|
| 102 |
+
#include <ATen/ops/absolute_compositeimplicitautograd_dispatch.h>
|
| 103 |
+
#include <ATen/ops/adaptive_avg_pool1d_compositeimplicitautograd_dispatch.h>
|
| 104 |
+
#include <ATen/ops/adaptive_avg_pool2d_compositeimplicitautograd_dispatch.h>
|
| 105 |
+
#include <ATen/ops/adaptive_avg_pool3d_compositeimplicitautograd_dispatch.h>
|
| 106 |
+
#include <ATen/ops/adaptive_max_pool1d_compositeimplicitautograd_dispatch.h>
|
| 107 |
+
#include <ATen/ops/adjoint_compositeimplicitautograd_dispatch.h>
|
| 108 |
+
#include <ATen/ops/affine_grid_generator_backward_compositeimplicitautograd_dispatch.h>
|
| 109 |
+
#include <ATen/ops/align_as_compositeimplicitautograd_dispatch.h>
|
| 110 |
+
#include <ATen/ops/align_tensors_compositeimplicitautograd_dispatch.h>
|
| 111 |
+
#include <ATen/ops/align_to_compositeimplicitautograd_dispatch.h>
|
| 112 |
+
#include <ATen/ops/all_compositeimplicitautograd_dispatch.h>
|
| 113 |
+
#include <ATen/ops/alpha_dropout_compositeimplicitautograd_dispatch.h>
|
| 114 |
+
#include <ATen/ops/and_compositeimplicitautograd_dispatch.h>
|
| 115 |
+
#include <ATen/ops/any_compositeimplicitautograd_dispatch.h>
|
| 116 |
+
#include <ATen/ops/arccos_compositeimplicitautograd_dispatch.h>
|
| 117 |
+
#include <ATen/ops/arccosh_compositeimplicitautograd_dispatch.h>
|
| 118 |
+
#include <ATen/ops/arcsin_compositeimplicitautograd_dispatch.h>
|
| 119 |
+
#include <ATen/ops/arcsinh_compositeimplicitautograd_dispatch.h>
|
| 120 |
+
#include <ATen/ops/arctan_compositeimplicitautograd_dispatch.h>
|
| 121 |
+
#include <ATen/ops/arctan2_compositeimplicitautograd_dispatch.h>
|
| 122 |
+
#include <ATen/ops/arctanh_compositeimplicitautograd_dispatch.h>
|
| 123 |
+
#include <ATen/ops/argsort_compositeimplicitautograd_dispatch.h>
|
| 124 |
+
#include <ATen/ops/argwhere_compositeimplicitautograd_dispatch.h>
|
| 125 |
+
#include <ATen/ops/atleast_1d_compositeimplicitautograd_dispatch.h>
|
| 126 |
+
#include <ATen/ops/atleast_2d_compositeimplicitautograd_dispatch.h>
|
| 127 |
+
#include <ATen/ops/atleast_3d_compositeimplicitautograd_dispatch.h>
|
| 128 |
+
#include <ATen/ops/avg_pool1d_compositeimplicitautograd_dispatch.h>
|
| 129 |
+
#include <ATen/ops/batch_norm_compositeimplicitautograd_dispatch.h>
|
| 130 |
+
#include <ATen/ops/bilinear_compositeimplicitautograd_dispatch.h>
|
| 131 |
+
#include <ATen/ops/broadcast_tensors_compositeimplicitautograd_dispatch.h>
|
| 132 |
+
#include <ATen/ops/broadcast_to_compositeimplicitautograd_dispatch.h>
|
| 133 |
+
#include <ATen/ops/can_cast_compositeimplicitautograd_dispatch.h>
|
| 134 |
+
#include <ATen/ops/cartesian_prod_compositeimplicitautograd_dispatch.h>
|
| 135 |
+
#include <ATen/ops/cat_compositeimplicitautograd_dispatch.h>
|
| 136 |
+
#include <ATen/ops/cdist_compositeimplicitautograd_dispatch.h>
|
| 137 |
+
#include <ATen/ops/chain_matmul_compositeimplicitautograd_dispatch.h>
|
| 138 |
+
#include <ATen/ops/chalf_compositeimplicitautograd_dispatch.h>
|
| 139 |
+
#include <ATen/ops/choose_qparams_optimized_compositeimplicitautograd_dispatch.h>
|
| 140 |
+
#include <ATen/ops/chunk_compositeimplicitautograd_dispatch.h>
|
| 141 |
+
#include <ATen/ops/clip_compositeimplicitautograd_dispatch.h>
|
| 142 |
+
#include <ATen/ops/coalesce_compositeimplicitautograd_dispatch.h>
|
| 143 |
+
#include <ATen/ops/column_stack_compositeimplicitautograd_dispatch.h>
|
| 144 |
+
#include <ATen/ops/combinations_compositeimplicitautograd_dispatch.h>
|
| 145 |
+
#include <ATen/ops/concat_compositeimplicitautograd_dispatch.h>
|
| 146 |
+
#include <ATen/ops/concatenate_compositeimplicitautograd_dispatch.h>
|
| 147 |
+
#include <ATen/ops/conj_compositeimplicitautograd_dispatch.h>
|
| 148 |
+
#include <ATen/ops/conj_physical_compositeimplicitautograd_dispatch.h>
|
| 149 |
+
#include <ATen/ops/contiguous_compositeimplicitautograd_dispatch.h>
|
| 150 |
+
#include <ATen/ops/conv1d_compositeimplicitautograd_dispatch.h>
|
| 151 |
+
#include <ATen/ops/conv2d_compositeimplicitautograd_dispatch.h>
|
| 152 |
+
#include <ATen/ops/conv3d_compositeimplicitautograd_dispatch.h>
|
| 153 |
+
#include <ATen/ops/conv_tbc_backward_compositeimplicitautograd_dispatch.h>
|
| 154 |
+
#include <ATen/ops/conv_transpose1d_compositeimplicitautograd_dispatch.h>
|
| 155 |
+
#include <ATen/ops/conv_transpose2d_compositeimplicitautograd_dispatch.h>
|
| 156 |
+
#include <ATen/ops/conv_transpose3d_compositeimplicitautograd_dispatch.h>
|
| 157 |
+
#include <ATen/ops/corrcoef_compositeimplicitautograd_dispatch.h>
|
| 158 |
+
#include <ATen/ops/cosine_embedding_loss_compositeimplicitautograd_dispatch.h>
|
| 159 |
+
#include <ATen/ops/cosine_similarity_compositeimplicitautograd_dispatch.h>
|
| 160 |
+
#include <ATen/ops/cov_compositeimplicitautograd_dispatch.h>
|
| 161 |
+
#include <ATen/ops/cross_compositeimplicitautograd_dispatch.h>
|
| 162 |
+
#include <ATen/ops/cross_entropy_loss_compositeimplicitautograd_dispatch.h>
|
| 163 |
+
#include <ATen/ops/ctc_loss_compositeimplicitautograd_dispatch.h>
|
| 164 |
+
#include <ATen/ops/cudnn_is_acceptable_compositeimplicitautograd_dispatch.h>
|
| 165 |
+
#include <ATen/ops/cummax_compositeimplicitautograd_dispatch.h>
|
| 166 |
+
#include <ATen/ops/cummaxmin_backward_compositeimplicitautograd_dispatch.h>
|
| 167 |
+
#include <ATen/ops/cummin_compositeimplicitautograd_dispatch.h>
|
| 168 |
+
#include <ATen/ops/cumprod_compositeimplicitautograd_dispatch.h>
|
| 169 |
+
#include <ATen/ops/cumprod_backward_compositeimplicitautograd_dispatch.h>
|
| 170 |
+
#include <ATen/ops/cumsum_compositeimplicitautograd_dispatch.h>
|
| 171 |
+
#include <ATen/ops/cumulative_trapezoid_compositeimplicitautograd_dispatch.h>
|
| 172 |
+
#include <ATen/ops/data_compositeimplicitautograd_dispatch.h>
|
| 173 |
+
#include <ATen/ops/det_compositeimplicitautograd_dispatch.h>
|
| 174 |
+
#include <ATen/ops/diag_compositeimplicitautograd_dispatch.h>
|
| 175 |
+
#include <ATen/ops/diagflat_compositeimplicitautograd_dispatch.h>
|
| 176 |
+
#include <ATen/ops/diagonal_compositeimplicitautograd_dispatch.h>
|
| 177 |
+
#include <ATen/ops/diff_compositeimplicitautograd_dispatch.h>
|
| 178 |
+
#include <ATen/ops/divide_compositeimplicitautograd_dispatch.h>
|
| 179 |
+
#include <ATen/ops/dropout_compositeimplicitautograd_dispatch.h>
|
| 180 |
+
#include <ATen/ops/dsplit_compositeimplicitautograd_dispatch.h>
|
| 181 |
+
#include <ATen/ops/dstack_compositeimplicitautograd_dispatch.h>
|
| 182 |
+
#include <ATen/ops/einsum_compositeimplicitautograd_dispatch.h>
|
| 183 |
+
#include <ATen/ops/embedding_backward_compositeimplicitautograd_dispatch.h>
|
| 184 |
+
#include <ATen/ops/embedding_bag_compositeimplicitautograd_dispatch.h>
|
| 185 |
+
#include <ATen/ops/embedding_sparse_backward_compositeimplicitautograd_dispatch.h>
|
| 186 |
+
#include <ATen/ops/empty_compositeimplicitautograd_dispatch.h>
|
| 187 |
+
#include <ATen/ops/expand_as_compositeimplicitautograd_dispatch.h>
|
| 188 |
+
#include <ATen/ops/fake_quantize_per_channel_affine_compositeimplicitautograd_dispatch.h>
|
| 189 |
+
#include <ATen/ops/fake_quantize_per_channel_affine_cachemask_backward_compositeimplicitautograd_dispatch.h>
|
| 190 |
+
#include <ATen/ops/fake_quantize_per_tensor_affine_compositeimplicitautograd_dispatch.h>
|
| 191 |
+
#include <ATen/ops/fake_quantize_per_tensor_affine_cachemask_backward_compositeimplicitautograd_dispatch.h>
|
| 192 |
+
#include <ATen/ops/fbgemm_linear_fp16_weight_compositeimplicitautograd_dispatch.h>
|
| 193 |
+
#include <ATen/ops/fbgemm_linear_fp16_weight_fp32_activation_compositeimplicitautograd_dispatch.h>
|
| 194 |
+
#include <ATen/ops/fbgemm_linear_int8_weight_compositeimplicitautograd_dispatch.h>
|
| 195 |
+
#include <ATen/ops/fbgemm_linear_int8_weight_fp32_activation_compositeimplicitautograd_dispatch.h>
|
| 196 |
+
#include <ATen/ops/fbgemm_linear_quantize_weight_compositeimplicitautograd_dispatch.h>
|
| 197 |
+
#include <ATen/ops/fbgemm_pack_gemm_matrix_fp16_compositeimplicitautograd_dispatch.h>
|
| 198 |
+
#include <ATen/ops/fbgemm_pack_quantized_matrix_compositeimplicitautograd_dispatch.h>
|
| 199 |
+
#include <ATen/ops/feature_alpha_dropout_compositeimplicitautograd_dispatch.h>
|
| 200 |
+
#include <ATen/ops/feature_dropout_compositeimplicitautograd_dispatch.h>
|
| 201 |
+
#include <ATen/ops/fft_fft_compositeimplicitautograd_dispatch.h>
|
| 202 |
+
#include <ATen/ops/fft_fft2_compositeimplicitautograd_dispatch.h>
|
| 203 |
+
#include <ATen/ops/fft_fftn_compositeimplicitautograd_dispatch.h>
|
| 204 |
+
#include <ATen/ops/fft_fftshift_compositeimplicitautograd_dispatch.h>
|
| 205 |
+
#include <ATen/ops/fft_hfft_compositeimplicitautograd_dispatch.h>
|
| 206 |
+
#include <ATen/ops/fft_hfft2_compositeimplicitautograd_dispatch.h>
|
| 207 |
+
#include <ATen/ops/fft_hfftn_compositeimplicitautograd_dispatch.h>
|
| 208 |
+
#include <ATen/ops/fft_ifft_compositeimplicitautograd_dispatch.h>
|
| 209 |
+
#include <ATen/ops/fft_ifft2_compositeimplicitautograd_dispatch.h>
|
| 210 |
+
#include <ATen/ops/fft_ifftn_compositeimplicitautograd_dispatch.h>
|
| 211 |
+
#include <ATen/ops/fft_ifftshift_compositeimplicitautograd_dispatch.h>
|
| 212 |
+
#include <ATen/ops/fft_ihfft_compositeimplicitautograd_dispatch.h>
|
| 213 |
+
#include <ATen/ops/fft_ihfft2_compositeimplicitautograd_dispatch.h>
|
| 214 |
+
#include <ATen/ops/fft_ihfftn_compositeimplicitautograd_dispatch.h>
|
| 215 |
+
#include <ATen/ops/fft_irfft_compositeimplicitautograd_dispatch.h>
|
| 216 |
+
#include <ATen/ops/fft_irfft2_compositeimplicitautograd_dispatch.h>
|
| 217 |
+
#include <ATen/ops/fft_irfftn_compositeimplicitautograd_dispatch.h>
|
| 218 |
+
#include <ATen/ops/fft_rfft_compositeimplicitautograd_dispatch.h>
|
| 219 |
+
#include <ATen/ops/fft_rfft2_compositeimplicitautograd_dispatch.h>
|
| 220 |
+
#include <ATen/ops/fft_rfftn_compositeimplicitautograd_dispatch.h>
|
| 221 |
+
#include <ATen/ops/fill_diagonal_compositeimplicitautograd_dispatch.h>
|
| 222 |
+
#include <ATen/ops/fix_compositeimplicitautograd_dispatch.h>
|
| 223 |
+
#include <ATen/ops/flatten_compositeimplicitautograd_dispatch.h>
|
| 224 |
+
#include <ATen/ops/flatten_dense_tensors_compositeimplicitautograd_dispatch.h>
|
| 225 |
+
#include <ATen/ops/fliplr_compositeimplicitautograd_dispatch.h>
|
| 226 |
+
#include <ATen/ops/flipud_compositeimplicitautograd_dispatch.h>
|
| 227 |
+
#include <ATen/ops/float_power_compositeimplicitautograd_dispatch.h>
|
| 228 |
+
#include <ATen/ops/frobenius_norm_compositeimplicitautograd_dispatch.h>
|
| 229 |
+
#include <ATen/ops/fused_moving_avg_obs_fake_quant_compositeimplicitautograd_dispatch.h>
|
| 230 |
+
#include <ATen/ops/gather_compositeimplicitautograd_dispatch.h>
|
| 231 |
+
#include <ATen/ops/gather_backward_compositeimplicitautograd_dispatch.h>
|
| 232 |
+
#include <ATen/ops/ger_compositeimplicitautograd_dispatch.h>
|
| 233 |
+
#include <ATen/ops/gradient_compositeimplicitautograd_dispatch.h>
|
| 234 |
+
#include <ATen/ops/greater_compositeimplicitautograd_dispatch.h>
|
| 235 |
+
#include <ATen/ops/greater_equal_compositeimplicitautograd_dispatch.h>
|
| 236 |
+
#include <ATen/ops/grid_sampler_compositeimplicitautograd_dispatch.h>
|
| 237 |
+
#include <ATen/ops/group_norm_compositeimplicitautograd_dispatch.h>
|
| 238 |
+
#include <ATen/ops/gru_compositeimplicitautograd_dispatch.h>
|
| 239 |
+
#include <ATen/ops/gru_cell_compositeimplicitautograd_dispatch.h>
|
| 240 |
+
#include <ATen/ops/hinge_embedding_loss_compositeimplicitautograd_dispatch.h>
|
| 241 |
+
#include <ATen/ops/histogramdd_compositeimplicitautograd_dispatch.h>
|
| 242 |
+
#include <ATen/ops/hsplit_compositeimplicitautograd_dispatch.h>
|
| 243 |
+
#include <ATen/ops/hstack_compositeimplicitautograd_dispatch.h>
|
| 244 |
+
#include <ATen/ops/imag_compositeimplicitautograd_dispatch.h>
|
| 245 |
+
#include <ATen/ops/index_add_compositeimplicitautograd_dispatch.h>
|
| 246 |
+
#include <ATen/ops/index_copy_compositeimplicitautograd_dispatch.h>
|
| 247 |
+
#include <ATen/ops/index_fill_compositeimplicitautograd_dispatch.h>
|
| 248 |
+
#include <ATen/ops/index_select_compositeimplicitautograd_dispatch.h>
|
| 249 |
+
#include <ATen/ops/index_select_backward_compositeimplicitautograd_dispatch.h>
|
| 250 |
+
#include <ATen/ops/infinitely_differentiable_gelu_backward_compositeimplicitautograd_dispatch.h>
|
| 251 |
+
#include <ATen/ops/inner_compositeimplicitautograd_dispatch.h>
|
| 252 |
+
#include <ATen/ops/instance_norm_compositeimplicitautograd_dispatch.h>
|
| 253 |
+
#include <ATen/ops/inverse_compositeimplicitautograd_dispatch.h>
|
| 254 |
+
#include <ATen/ops/is_complex_compositeimplicitautograd_dispatch.h>
|
| 255 |
+
#include <ATen/ops/is_conj_compositeimplicitautograd_dispatch.h>
|
| 256 |
+
#include <ATen/ops/is_distributed_compositeimplicitautograd_dispatch.h>
|
| 257 |
+
#include <ATen/ops/is_floating_point_compositeimplicitautograd_dispatch.h>
|
| 258 |
+
#include <ATen/ops/is_inference_compositeimplicitautograd_dispatch.h>
|
| 259 |
+
#include <ATen/ops/is_leaf_compositeimplicitautograd_dispatch.h>
|
| 260 |
+
#include <ATen/ops/is_neg_compositeimplicitautograd_dispatch.h>
|
| 261 |
+
#include <ATen/ops/is_nonzero_compositeimplicitautograd_dispatch.h>
|
| 262 |
+
#include <ATen/ops/is_signed_compositeimplicitautograd_dispatch.h>
|
| 263 |
+
#include <ATen/ops/is_vulkan_available_compositeimplicitautograd_dispatch.h>
|
| 264 |
+
#include <ATen/ops/isclose_compositeimplicitautograd_dispatch.h>
|
| 265 |
+
#include <ATen/ops/isfinite_compositeimplicitautograd_dispatch.h>
|
| 266 |
+
#include <ATen/ops/isreal_compositeimplicitautograd_dispatch.h>
|
| 267 |
+
#include <ATen/ops/istft_compositeimplicitautograd_dispatch.h>
|
| 268 |
+
#include <ATen/ops/item_compositeimplicitautograd_dispatch.h>
|
| 269 |
+
#include <ATen/ops/kl_div_compositeimplicitautograd_dispatch.h>
|
| 270 |
+
#include <ATen/ops/kron_compositeimplicitautograd_dispatch.h>
|
| 271 |
+
#include <ATen/ops/kthvalue_compositeimplicitautograd_dispatch.h>
|
| 272 |
+
#include <ATen/ops/l1_loss_compositeimplicitautograd_dispatch.h>
|
| 273 |
+
#include <ATen/ops/layer_norm_compositeimplicitautograd_dispatch.h>
|
| 274 |
+
#include <ATen/ops/ldexp_compositeimplicitautograd_dispatch.h>
|
| 275 |
+
#include <ATen/ops/less_compositeimplicitautograd_dispatch.h>
|
| 276 |
+
#include <ATen/ops/less_equal_compositeimplicitautograd_dispatch.h>
|
| 277 |
+
#include <ATen/ops/linalg_cholesky_compositeimplicitautograd_dispatch.h>
|
| 278 |
+
#include <ATen/ops/linalg_cond_compositeimplicitautograd_dispatch.h>
|
| 279 |
+
#include <ATen/ops/linalg_det_compositeimplicitautograd_dispatch.h>
|
| 280 |
+
#include <ATen/ops/linalg_diagonal_compositeimplicitautograd_dispatch.h>
|
| 281 |
+
#include <ATen/ops/linalg_eigh_compositeimplicitautograd_dispatch.h>
|
| 282 |
+
#include <ATen/ops/linalg_eigvals_compositeimplicitautograd_dispatch.h>
|
| 283 |
+
#include <ATen/ops/linalg_eigvalsh_compositeimplicitautograd_dispatch.h>
|
| 284 |
+
#include <ATen/ops/linalg_inv_compositeimplicitautograd_dispatch.h>
|
| 285 |
+
#include <ATen/ops/linalg_ldl_factor_compositeimplicitautograd_dispatch.h>
|
| 286 |
+
#include <ATen/ops/linalg_lu_factor_compositeimplicitautograd_dispatch.h>
|
| 287 |
+
#include <ATen/ops/linalg_matmul_compositeimplicitautograd_dispatch.h>
|
| 288 |
+
#include <ATen/ops/linalg_matrix_norm_compositeimplicitautograd_dispatch.h>
|
| 289 |
+
#include <ATen/ops/linalg_matrix_power_compositeimplicitautograd_dispatch.h>
|
| 290 |
+
#include <ATen/ops/linalg_matrix_rank_compositeimplicitautograd_dispatch.h>
|
| 291 |
+
#include <ATen/ops/linalg_multi_dot_compositeimplicitautograd_dispatch.h>
|
| 292 |
+
#include <ATen/ops/linalg_norm_compositeimplicitautograd_dispatch.h>
|
| 293 |
+
#include <ATen/ops/linalg_pinv_compositeimplicitautograd_dispatch.h>
|
| 294 |
+
#include <ATen/ops/linalg_slogdet_compositeimplicitautograd_dispatch.h>
|
| 295 |
+
#include <ATen/ops/linalg_solve_compositeimplicitautograd_dispatch.h>
|
| 296 |
+
#include <ATen/ops/linalg_solve_ex_compositeimplicitautograd_dispatch.h>
|
| 297 |
+
#include <ATen/ops/linalg_svd_compositeimplicitautograd_dispatch.h>
|
| 298 |
+
#include <ATen/ops/linalg_svdvals_compositeimplicitautograd_dispatch.h>
|
| 299 |
+
#include <ATen/ops/linalg_tensorinv_compositeimplicitautograd_dispatch.h>
|
| 300 |
+
#include <ATen/ops/linalg_tensorsolve_compositeimplicitautograd_dispatch.h>
|
| 301 |
+
#include <ATen/ops/linalg_vander_compositeimplicitautograd_dispatch.h>
|
| 302 |
+
#include <ATen/ops/linalg_vecdot_compositeimplicitautograd_dispatch.h>
|
| 303 |
+
#include <ATen/ops/linear_compositeimplicitautograd_dispatch.h>
|
| 304 |
+
#include <ATen/ops/log_sigmoid_compositeimplicitautograd_dispatch.h>
|
| 305 |
+
#include <ATen/ops/log_softmax_compositeimplicitautograd_dispatch.h>
|
| 306 |
+
#include <ATen/ops/logcumsumexp_compositeimplicitautograd_dispatch.h>
|
| 307 |
+
#include <ATen/ops/logdet_compositeimplicitautograd_dispatch.h>
|
| 308 |
+
#include <ATen/ops/logsumexp_compositeimplicitautograd_dispatch.h>
|
| 309 |
+
#include <ATen/ops/lstm_compositeimplicitautograd_dispatch.h>
|
| 310 |
+
#include <ATen/ops/lstm_cell_compositeimplicitautograd_dispatch.h>
|
| 311 |
+
#include <ATen/ops/lu_solve_compositeimplicitautograd_dispatch.h>
|
| 312 |
+
#include <ATen/ops/mH_compositeimplicitautograd_dispatch.h>
|
| 313 |
+
#include <ATen/ops/mT_compositeimplicitautograd_dispatch.h>
|
| 314 |
+
#include <ATen/ops/margin_ranking_loss_compositeimplicitautograd_dispatch.h>
|
| 315 |
+
#include <ATen/ops/masked_select_backward_compositeimplicitautograd_dispatch.h>
|
| 316 |
+
#include <ATen/ops/matmul_compositeimplicitautograd_dispatch.h>
|
| 317 |
+
#include <ATen/ops/matrix_H_compositeimplicitautograd_dispatch.h>
|
| 318 |
+
#include <ATen/ops/matrix_exp_compositeimplicitautograd_dispatch.h>
|
| 319 |
+
#include <ATen/ops/matrix_exp_backward_compositeimplicitautograd_dispatch.h>
|
| 320 |
+
#include <ATen/ops/matrix_power_compositeimplicitautograd_dispatch.h>
|
| 321 |
+
#include <ATen/ops/max_compositeimplicitautograd_dispatch.h>
|
| 322 |
+
#include <ATen/ops/max_pool1d_compositeimplicitautograd_dispatch.h>
|
| 323 |
+
#include <ATen/ops/max_pool1d_with_indices_compositeimplicitautograd_dispatch.h>
|
| 324 |
+
#include <ATen/ops/max_pool2d_compositeimplicitautograd_dispatch.h>
|
| 325 |
+
#include <ATen/ops/max_pool3d_compositeimplicitautograd_dispatch.h>
|
| 326 |
+
#include <ATen/ops/mean_compositeimplicitautograd_dispatch.h>
|
| 327 |
+
#include <ATen/ops/median_compositeimplicitautograd_dispatch.h>
|
| 328 |
+
#include <ATen/ops/meshgrid_compositeimplicitautograd_dispatch.h>
|
| 329 |
+
#include <ATen/ops/min_compositeimplicitautograd_dispatch.h>
|
| 330 |
+
#include <ATen/ops/mish_backward_compositeimplicitautograd_dispatch.h>
|
| 331 |
+
#include <ATen/ops/mode_compositeimplicitautograd_dispatch.h>
|
| 332 |
+
#include <ATen/ops/moveaxis_compositeimplicitautograd_dispatch.h>
|
| 333 |
+
#include <ATen/ops/movedim_compositeimplicitautograd_dispatch.h>
|
| 334 |
+
#include <ATen/ops/msort_compositeimplicitautograd_dispatch.h>
|
| 335 |
+
#include <ATen/ops/multilabel_margin_loss_compositeimplicitautograd_dispatch.h>
|
| 336 |
+
#include <ATen/ops/multiply_compositeimplicitautograd_dispatch.h>
|
| 337 |
+
#include <ATen/ops/nanmean_compositeimplicitautograd_dispatch.h>
|
| 338 |
+
#include <ATen/ops/nanmedian_compositeimplicitautograd_dispatch.h>
|
| 339 |
+
#include <ATen/ops/nanquantile_compositeimplicitautograd_dispatch.h>
|
| 340 |
+
#include <ATen/ops/narrow_compositeimplicitautograd_dispatch.h>
|
| 341 |
+
#include <ATen/ops/native_channel_shuffle_compositeimplicitautograd_dispatch.h>
|
| 342 |
+
#include <ATen/ops/negative_compositeimplicitautograd_dispatch.h>
|
| 343 |
+
#include <ATen/ops/nested_to_padded_tensor_compositeimplicitautograd_dispatch.h>
|
| 344 |
+
#include <ATen/ops/nll_loss_compositeimplicitautograd_dispatch.h>
|
| 345 |
+
#include <ATen/ops/nll_loss2d_compositeimplicitautograd_dispatch.h>
|
| 346 |
+
#include <ATen/ops/nll_loss_nd_compositeimplicitautograd_dispatch.h>
|
| 347 |
+
#include <ATen/ops/nonzero_numpy_compositeimplicitautograd_dispatch.h>
|
| 348 |
+
#include <ATen/ops/norm_compositeimplicitautograd_dispatch.h>
|
| 349 |
+
#include <ATen/ops/norm_except_dim_compositeimplicitautograd_dispatch.h>
|
| 350 |
+
#include <ATen/ops/not_equal_compositeimplicitautograd_dispatch.h>
|
| 351 |
+
#include <ATen/ops/nuclear_norm_compositeimplicitautograd_dispatch.h>
|
| 352 |
+
#include <ATen/ops/numpy_T_compositeimplicitautograd_dispatch.h>
|
| 353 |
+
#include <ATen/ops/one_hot_compositeimplicitautograd_dispatch.h>
|
| 354 |
+
#include <ATen/ops/or_compositeimplicitautograd_dispatch.h>
|
| 355 |
+
#include <ATen/ops/orgqr_compositeimplicitautograd_dispatch.h>
|
| 356 |
+
#include <ATen/ops/outer_compositeimplicitautograd_dispatch.h>
|
| 357 |
+
#include <ATen/ops/output_nr_compositeimplicitautograd_dispatch.h>
|
| 358 |
+
#include <ATen/ops/pad_compositeimplicitautograd_dispatch.h>
|
| 359 |
+
#include <ATen/ops/pad_sequence_compositeimplicitautograd_dispatch.h>
|
| 360 |
+
#include <ATen/ops/pairwise_distance_compositeimplicitautograd_dispatch.h>
|
| 361 |
+
#include <ATen/ops/pdist_compositeimplicitautograd_dispatch.h>
|
| 362 |
+
#include <ATen/ops/pin_memory_compositeimplicitautograd_dispatch.h>
|
| 363 |
+
#include <ATen/ops/pinverse_compositeimplicitautograd_dispatch.h>
|
| 364 |
+
#include <ATen/ops/poisson_nll_loss_compositeimplicitautograd_dispatch.h>
|
| 365 |
+
#include <ATen/ops/positive_compositeimplicitautograd_dispatch.h>
|
| 366 |
+
#include <ATen/ops/prelu_compositeimplicitautograd_dispatch.h>
|
| 367 |
+
#include <ATen/ops/prod_compositeimplicitautograd_dispatch.h>
|
| 368 |
+
#include <ATen/ops/promote_types_compositeimplicitautograd_dispatch.h>
|
| 369 |
+
#include <ATen/ops/qr_compositeimplicitautograd_dispatch.h>
|
| 370 |
+
#include <ATen/ops/quantile_compositeimplicitautograd_dispatch.h>
|
| 371 |
+
#include <ATen/ops/quantized_gru_cell_compositeimplicitautograd_dispatch.h>
|
| 372 |
+
#include <ATen/ops/quantized_lstm_cell_compositeimplicitautograd_dispatch.h>
|
| 373 |
+
#include <ATen/ops/quantized_rnn_relu_cell_compositeimplicitautograd_dispatch.h>
|
| 374 |
+
#include <ATen/ops/quantized_rnn_tanh_cell_compositeimplicitautograd_dispatch.h>
|
| 375 |
+
#include <ATen/ops/rand_compositeimplicitautograd_dispatch.h>
|
| 376 |
+
#include <ATen/ops/randn_compositeimplicitautograd_dispatch.h>
|
| 377 |
+
#include <ATen/ops/ravel_compositeimplicitautograd_dispatch.h>
|
| 378 |
+
#include <ATen/ops/real_compositeimplicitautograd_dispatch.h>
|
| 379 |
+
#include <ATen/ops/refine_names_compositeimplicitautograd_dispatch.h>
|
| 380 |
+
#include <ATen/ops/relu6_compositeimplicitautograd_dispatch.h>
|
| 381 |
+
#include <ATen/ops/rename_compositeimplicitautograd_dispatch.h>
|
| 382 |
+
#include <ATen/ops/repeat_interleave_compositeimplicitautograd_dispatch.h>
|
| 383 |
+
#include <ATen/ops/requires_grad_compositeimplicitautograd_dispatch.h>
|
| 384 |
+
#include <ATen/ops/reshape_compositeimplicitautograd_dispatch.h>
|
| 385 |
+
#include <ATen/ops/reshape_as_compositeimplicitautograd_dispatch.h>
|
| 386 |
+
#include <ATen/ops/resolve_conj_compositeimplicitautograd_dispatch.h>
|
| 387 |
+
#include <ATen/ops/resolve_neg_compositeimplicitautograd_dispatch.h>
|
| 388 |
+
#include <ATen/ops/result_type_compositeimplicitautograd_dispatch.h>
|
| 389 |
+
#include <ATen/ops/retain_grad_compositeimplicitautograd_dispatch.h>
|
| 390 |
+
#include <ATen/ops/retains_grad_compositeimplicitautograd_dispatch.h>
|
| 391 |
+
#include <ATen/ops/rms_norm_compositeimplicitautograd_dispatch.h>
|
| 392 |
+
#include <ATen/ops/rnn_relu_compositeimplicitautograd_dispatch.h>
|
| 393 |
+
#include <ATen/ops/rnn_relu_cell_compositeimplicitautograd_dispatch.h>
|
| 394 |
+
#include <ATen/ops/rnn_tanh_compositeimplicitautograd_dispatch.h>
|
| 395 |
+
#include <ATen/ops/rnn_tanh_cell_compositeimplicitautograd_dispatch.h>
|
| 396 |
+
#include <ATen/ops/row_stack_compositeimplicitautograd_dispatch.h>
|
| 397 |
+
#include <ATen/ops/rrelu_compositeimplicitautograd_dispatch.h>
|
| 398 |
+
#include <ATen/ops/scaled_dot_product_attention_compositeimplicitautograd_dispatch.h>
|
| 399 |
+
#include <ATen/ops/scatter_compositeimplicitautograd_dispatch.h>
|
| 400 |
+
#include <ATen/ops/scatter_add_compositeimplicitautograd_dispatch.h>
|
| 401 |
+
#include <ATen/ops/select_compositeimplicitautograd_dispatch.h>
|
| 402 |
+
#include <ATen/ops/selu_compositeimplicitautograd_dispatch.h>
|
| 403 |
+
#include <ATen/ops/set_compositeimplicitautograd_dispatch.h>
|
| 404 |
+
#include <ATen/ops/set_data_compositeimplicitautograd_dispatch.h>
|
| 405 |
+
#include <ATen/ops/silu_backward_compositeimplicitautograd_dispatch.h>
|
| 406 |
+
#include <ATen/ops/size_compositeimplicitautograd_dispatch.h>
|
| 407 |
+
#include <ATen/ops/slogdet_compositeimplicitautograd_dispatch.h>
|
| 408 |
+
#include <ATen/ops/slow_conv3d_compositeimplicitautograd_dispatch.h>
|
| 409 |
+
#include <ATen/ops/smm_compositeimplicitautograd_dispatch.h>
|
| 410 |
+
#include <ATen/ops/softmax_compositeimplicitautograd_dispatch.h>
|
| 411 |
+
#include <ATen/ops/sort_compositeimplicitautograd_dispatch.h>
|
| 412 |
+
#include <ATen/ops/sparse_bsc_tensor_compositeimplicitautograd_dispatch.h>
|
| 413 |
+
#include <ATen/ops/sparse_bsr_tensor_compositeimplicitautograd_dispatch.h>
|
| 414 |
+
#include <ATen/ops/sparse_coo_tensor_compositeimplicitautograd_dispatch.h>
|
| 415 |
+
#include <ATen/ops/sparse_csc_tensor_compositeimplicitautograd_dispatch.h>
|
| 416 |
+
#include <ATen/ops/sparse_csr_tensor_compositeimplicitautograd_dispatch.h>
|
| 417 |
+
#include <ATen/ops/special_digamma_compositeimplicitautograd_dispatch.h>
|
| 418 |
+
#include <ATen/ops/special_erf_compositeimplicitautograd_dispatch.h>
|
| 419 |
+
#include <ATen/ops/special_erfc_compositeimplicitautograd_dispatch.h>
|
| 420 |
+
#include <ATen/ops/special_erfinv_compositeimplicitautograd_dispatch.h>
|
| 421 |
+
#include <ATen/ops/special_exp2_compositeimplicitautograd_dispatch.h>
|
| 422 |
+
#include <ATen/ops/special_expit_compositeimplicitautograd_dispatch.h>
|
| 423 |
+
#include <ATen/ops/special_expm1_compositeimplicitautograd_dispatch.h>
|
| 424 |
+
#include <ATen/ops/special_gammainc_compositeimplicitautograd_dispatch.h>
|
| 425 |
+
#include <ATen/ops/special_gammaincc_compositeimplicitautograd_dispatch.h>
|
| 426 |
+
#include <ATen/ops/special_gammaln_compositeimplicitautograd_dispatch.h>
|
| 427 |
+
#include <ATen/ops/special_i0_compositeimplicitautograd_dispatch.h>
|
| 428 |
+
#include <ATen/ops/special_log1p_compositeimplicitautograd_dispatch.h>
|
| 429 |
+
#include <ATen/ops/special_log_softmax_compositeimplicitautograd_dispatch.h>
|
| 430 |
+
#include <ATen/ops/special_logit_compositeimplicitautograd_dispatch.h>
|
| 431 |
+
#include <ATen/ops/special_logsumexp_compositeimplicitautograd_dispatch.h>
|
| 432 |
+
#include <ATen/ops/special_multigammaln_compositeimplicitautograd_dispatch.h>
|
| 433 |
+
#include <ATen/ops/special_ndtr_compositeimplicitautograd_dispatch.h>
|
| 434 |
+
#include <ATen/ops/special_polygamma_compositeimplicitautograd_dispatch.h>
|
| 435 |
+
#include <ATen/ops/special_psi_compositeimplicitautograd_dispatch.h>
|
| 436 |
+
#include <ATen/ops/special_round_compositeimplicitautograd_dispatch.h>
|
| 437 |
+
#include <ATen/ops/special_sinc_compositeimplicitautograd_dispatch.h>
|
| 438 |
+
#include <ATen/ops/special_softmax_compositeimplicitautograd_dispatch.h>
|
| 439 |
+
#include <ATen/ops/special_xlogy_compositeimplicitautograd_dispatch.h>
|
| 440 |
+
#include <ATen/ops/split_compositeimplicitautograd_dispatch.h>
|
| 441 |
+
#include <ATen/ops/square_compositeimplicitautograd_dispatch.h>
|
| 442 |
+
#include <ATen/ops/squeeze_compositeimplicitautograd_dispatch.h>
|
| 443 |
+
#include <ATen/ops/sspaddmm_compositeimplicitautograd_dispatch.h>
|
| 444 |
+
#include <ATen/ops/std_compositeimplicitautograd_dispatch.h>
|
| 445 |
+
#include <ATen/ops/std_mean_compositeimplicitautograd_dispatch.h>
|
| 446 |
+
#include <ATen/ops/stft_compositeimplicitautograd_dispatch.h>
|
| 447 |
+
#include <ATen/ops/stride_compositeimplicitautograd_dispatch.h>
|
| 448 |
+
#include <ATen/ops/subtract_compositeimplicitautograd_dispatch.h>
|
| 449 |
+
#include <ATen/ops/sum_compositeimplicitautograd_dispatch.h>
|
| 450 |
+
#include <ATen/ops/sum_to_size_compositeimplicitautograd_dispatch.h>
|
| 451 |
+
#include <ATen/ops/svd_compositeimplicitautograd_dispatch.h>
|
| 452 |
+
#include <ATen/ops/swapaxes_compositeimplicitautograd_dispatch.h>
|
| 453 |
+
#include <ATen/ops/swapdims_compositeimplicitautograd_dispatch.h>
|
| 454 |
+
#include <ATen/ops/sym_numel_compositeimplicitautograd_dispatch.h>
|
| 455 |
+
#include <ATen/ops/sym_size_compositeimplicitautograd_dispatch.h>
|
| 456 |
+
#include <ATen/ops/sym_storage_offset_compositeimplicitautograd_dispatch.h>
|
| 457 |
+
#include <ATen/ops/sym_stride_compositeimplicitautograd_dispatch.h>
|
| 458 |
+
#include <ATen/ops/take_along_dim_compositeimplicitautograd_dispatch.h>
|
| 459 |
+
#include <ATen/ops/tensor_split_compositeimplicitautograd_dispatch.h>
|
| 460 |
+
#include <ATen/ops/tensordot_compositeimplicitautograd_dispatch.h>
|
| 461 |
+
#include <ATen/ops/thnn_conv2d_compositeimplicitautograd_dispatch.h>
|
| 462 |
+
#include <ATen/ops/tile_compositeimplicitautograd_dispatch.h>
|
| 463 |
+
#include <ATen/ops/to_compositeimplicitautograd_dispatch.h>
|
| 464 |
+
#include <ATen/ops/to_dense_compositeimplicitautograd_dispatch.h>
|
| 465 |
+
#include <ATen/ops/to_dense_backward_compositeimplicitautograd_dispatch.h>
|
| 466 |
+
#include <ATen/ops/to_mkldnn_backward_compositeimplicitautograd_dispatch.h>
|
| 467 |
+
#include <ATen/ops/to_sparse_compositeimplicitautograd_dispatch.h>
|
| 468 |
+
#include <ATen/ops/to_sparse_bsc_compositeimplicitautograd_dispatch.h>
|
| 469 |
+
#include <ATen/ops/to_sparse_bsr_compositeimplicitautograd_dispatch.h>
|
| 470 |
+
#include <ATen/ops/to_sparse_csc_compositeimplicitautograd_dispatch.h>
|
| 471 |
+
#include <ATen/ops/to_sparse_csr_compositeimplicitautograd_dispatch.h>
|
| 472 |
+
#include <ATen/ops/trace_backward_compositeimplicitautograd_dispatch.h>
|
| 473 |
+
#include <ATen/ops/transpose_compositeimplicitautograd_dispatch.h>
|
| 474 |
+
#include <ATen/ops/trapezoid_compositeimplicitautograd_dispatch.h>
|
| 475 |
+
#include <ATen/ops/trapz_compositeimplicitautograd_dispatch.h>
|
| 476 |
+
#include <ATen/ops/triplet_margin_loss_compositeimplicitautograd_dispatch.h>
|
| 477 |
+
#include <ATen/ops/true_divide_compositeimplicitautograd_dispatch.h>
|
| 478 |
+
#include <ATen/ops/type_as_compositeimplicitautograd_dispatch.h>
|
| 479 |
+
#include <ATen/ops/unbind_compositeimplicitautograd_dispatch.h>
|
| 480 |
+
#include <ATen/ops/unflatten_compositeimplicitautograd_dispatch.h>
|
| 481 |
+
#include <ATen/ops/unflatten_dense_tensors_compositeimplicitautograd_dispatch.h>
|
| 482 |
+
#include <ATen/ops/unsafe_chunk_compositeimplicitautograd_dispatch.h>
|
| 483 |
+
#include <ATen/ops/upsample_bicubic2d_compositeimplicitautograd_dispatch.h>
|
| 484 |
+
#include <ATen/ops/upsample_bilinear2d_compositeimplicitautograd_dispatch.h>
|
| 485 |
+
#include <ATen/ops/upsample_linear1d_compositeimplicitautograd_dispatch.h>
|
| 486 |
+
#include <ATen/ops/upsample_nearest1d_compositeimplicitautograd_dispatch.h>
|
| 487 |
+
#include <ATen/ops/upsample_nearest2d_compositeimplicitautograd_dispatch.h>
|
| 488 |
+
#include <ATen/ops/upsample_nearest3d_compositeimplicitautograd_dispatch.h>
|
| 489 |
+
#include <ATen/ops/upsample_trilinear3d_compositeimplicitautograd_dispatch.h>
|
| 490 |
+
#include <ATen/ops/value_selecting_reduction_backward_compositeimplicitautograd_dispatch.h>
|
| 491 |
+
#include <ATen/ops/vander_compositeimplicitautograd_dispatch.h>
|
| 492 |
+
#include <ATen/ops/var_compositeimplicitautograd_dispatch.h>
|
| 493 |
+
#include <ATen/ops/var_mean_compositeimplicitautograd_dispatch.h>
|
| 494 |
+
#include <ATen/ops/view_as_compositeimplicitautograd_dispatch.h>
|
| 495 |
+
#include <ATen/ops/vsplit_compositeimplicitautograd_dispatch.h>
|
| 496 |
+
#include <ATen/ops/vstack_compositeimplicitautograd_dispatch.h>
|
| 497 |
+
#include <ATen/ops/where_compositeimplicitautograd_dispatch.h>
|
| 498 |
+
#include <ATen/ops/xor_compositeimplicitautograd_dispatch.h>
|
| 499 |
+
|
| 500 |
+
|
| 501 |
+
|
lib/python3.12/site-packages/torch/include/ATen/Dispatch.h
ADDED
|
@@ -0,0 +1,808 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <c10/core/ScalarType.h>
|
| 4 |
+
#include <c10/macros/Macros.h>
|
| 5 |
+
#include <c10/util/Exception.h>
|
| 6 |
+
#include <c10/util/Half.h>
|
| 7 |
+
#include <c10/util/Metaprogramming.h>
|
| 8 |
+
#include <c10/util/complex.h>
|
| 9 |
+
#include <c10/util/string_view.h>
|
| 10 |
+
|
| 11 |
+
#ifdef __CUDACC__
|
| 12 |
+
#include <cuda.h> // For CUDA_VERSION
|
| 13 |
+
#endif
|
| 14 |
+
|
| 15 |
+
#ifdef TEMPLATE_SELECTIVE_BUILD
|
| 16 |
+
#include <ATen/selected_mobile_ops.h>
|
| 17 |
+
#else
|
| 18 |
+
namespace at {
|
| 19 |
+
/**
|
| 20 |
+
* The method should_include_kernel_dtype() returns true/false
|
| 21 |
+
* based on whether the switching code for a specific dtype should be
|
| 22 |
+
* included based on build time constants generated from tracing model
|
| 23 |
+
* execution. This method will be implemented via code-generation and
|
| 24 |
+
* included in this file when code-gen is ready.
|
| 25 |
+
*/
|
| 26 |
+
inline constexpr bool should_include_kernel_dtype(
|
| 27 |
+
const char* /*kernel_tag_str*/,
|
| 28 |
+
at::ScalarType /*scalar_type*/
|
| 29 |
+
) {
|
| 30 |
+
return true;
|
| 31 |
+
}
|
| 32 |
+
} // namespace at
|
| 33 |
+
#endif
|
| 34 |
+
|
| 35 |
+
/**
|
| 36 |
+
* In the Facebook internal build (using BUCK), this macro is enabled by
|
| 37 |
+
* passing in -c pt.enable_record_kernel_dtype=1 when building the tracer
|
| 38 |
+
* binary.
|
| 39 |
+
*/
|
| 40 |
+
#if defined ENABLE_RECORD_KERNEL_FUNCTION_DTYPE
|
| 41 |
+
namespace at::detail {
|
| 42 |
+
TORCH_API void record_kernel_function_dtype(std::string name);
|
| 43 |
+
} // namespace at::detail
|
| 44 |
+
|
| 45 |
+
#define RECORD_KERNEL_FUNCTION_DTYPE(NAME, enum_type) \
|
| 46 |
+
at::detail::record_kernel_function_dtype( \
|
| 47 |
+
std::string(NAME) + "$" + toString(enum_type));
|
| 48 |
+
#else
|
| 49 |
+
#define RECORD_KERNEL_FUNCTION_DTYPE(NAME, enum_type)
|
| 50 |
+
#endif
|
| 51 |
+
|
| 52 |
+
#define AT_PRIVATE_CHECK_SELECTIVE_BUILD(enum_type) \
|
| 53 |
+
do { \
|
| 54 |
+
if constexpr (!at::should_include_kernel_dtype( \
|
| 55 |
+
at_dispatch_name, enum_type)) { \
|
| 56 |
+
TORCH_CHECK( \
|
| 57 |
+
false, \
|
| 58 |
+
"dtype '", \
|
| 59 |
+
toString(enum_type), \
|
| 60 |
+
"' not selected for kernel tag ", \
|
| 61 |
+
at_dispatch_name); \
|
| 62 |
+
} \
|
| 63 |
+
} while (0)
|
| 64 |
+
|
| 65 |
+
#define AT_PRIVATE_CASE_TYPE_USING_HINT(enum_type, HINT, ...) \
|
| 66 |
+
case enum_type: { \
|
| 67 |
+
AT_PRIVATE_CHECK_SELECTIVE_BUILD(enum_type); \
|
| 68 |
+
using HINT [[maybe_unused]] = c10::impl::ScalarTypeToCPPTypeT<enum_type>; \
|
| 69 |
+
return __VA_ARGS__(); \
|
| 70 |
+
}
|
| 71 |
+
|
| 72 |
+
#define AT_DISPATCH_CASE(enum_type, ...) \
|
| 73 |
+
AT_PRIVATE_CASE_TYPE_USING_HINT(enum_type, scalar_t, __VA_ARGS__)
|
| 74 |
+
|
| 75 |
+
#define AT_DISPATCH_CASE_QINT(enum_type, scalar_type, ...) \
|
| 76 |
+
case enum_type: { \
|
| 77 |
+
AT_PRIVATE_CHECK_SELECTIVE_BUILD(enum_type); \
|
| 78 |
+
using scalar_t = scalar_type; \
|
| 79 |
+
using underlying_t [[maybe_unused]] = typename scalar_t::underlying; \
|
| 80 |
+
[[maybe_unused]] const auto& SCALAR_TYPE = enum_type; \
|
| 81 |
+
[[maybe_unused]] const auto& UNDERLYING_TYPE = toUnderlying(enum_type); \
|
| 82 |
+
return __VA_ARGS__(); \
|
| 83 |
+
}
|
| 84 |
+
|
| 85 |
+
#define AT_QINT_SUB_BYTE_PRIVATE_CASE_TYPE( \
|
| 86 |
+
enum_type, scalar_type, bitwidth, qmin, qmax, ...) \
|
| 87 |
+
case enum_type: { \
|
| 88 |
+
AT_PRIVATE_CHECK_SELECTIVE_BUILD(enum_type); \
|
| 89 |
+
using scalar_t = scalar_type; \
|
| 90 |
+
using underlying_t [[maybe_unused]] = typename scalar_t::underlying; \
|
| 91 |
+
[[maybe_unused]] const auto& SCALAR_TYPE = enum_type; \
|
| 92 |
+
[[maybe_unused]] const auto& UNDERLYING_TYPE = toUnderlying(enum_type); \
|
| 93 |
+
[[maybe_unused]] int bit_width = bitwidth; \
|
| 94 |
+
[[maybe_unused]] int64_t quant_min = qmin; \
|
| 95 |
+
[[maybe_unused]] int64_t quant_max = qmax; \
|
| 96 |
+
return __VA_ARGS__(); \
|
| 97 |
+
}
|
| 98 |
+
|
| 99 |
+
namespace detail {
|
| 100 |
+
|
| 101 |
+
inline at::ScalarType scalar_type(at::ScalarType s) {
|
| 102 |
+
return s;
|
| 103 |
+
}
|
| 104 |
+
|
| 105 |
+
} // namespace detail
|
| 106 |
+
|
| 107 |
+
// The AT_DISPATCH_* family of macros provides the ability to
|
| 108 |
+
// conveniently generate specializations of a kernel over all of the
|
| 109 |
+
// dtypes we care about in PyTorch. We call it "dispatch" because
|
| 110 |
+
// we are "dispatching" to the correct, dtype-specific kernel.
|
| 111 |
+
//
|
| 112 |
+
// A standard usage looks like:
|
| 113 |
+
//
|
| 114 |
+
// AT_DISPATCH_ALL_TYPES(self.scalar_type(), "op_name", [&] {
|
| 115 |
+
// // Your code here, with 'scalar_t' now defined to
|
| 116 |
+
// // be the dtype in question
|
| 117 |
+
// });
|
| 118 |
+
//
|
| 119 |
+
// There are many variations of this macro, so it's important to
|
| 120 |
+
// understand exactly /which/ dtypes you want to get instantiated, as
|
| 121 |
+
// well as what the "default" set is.
|
| 122 |
+
//
|
| 123 |
+
// The default set of dtypes that are instantiated (e.g., by
|
| 124 |
+
// AT_DISPATCH_ALL_TYPES) are floating point types (float, double),
|
| 125 |
+
// and integral types (int32_t, int64_t, int16_t, int8_t, uint8_t),
|
| 126 |
+
// but NOT booleans (bool), half-precision floats (Half) or
|
| 127 |
+
// complex number (c10::complex<float>, c10::complex<double>).
|
| 128 |
+
// This "cut" is somewhat historical (the default types are the
|
| 129 |
+
// ones that TH historically supported), but it also reflects the
|
| 130 |
+
// fact that the non-default types are "poorly" behaved (booleans
|
| 131 |
+
// are NOT integers mod 2, half precision operations ~essentially
|
| 132 |
+
// don't exist on CPU, complex numbers are an experimental application).
|
| 133 |
+
//
|
| 134 |
+
// Here are the questions you should generally ask to decide which
|
| 135 |
+
// dispatch you want:
|
| 136 |
+
//
|
| 137 |
+
// 1. Is this an integral or floating point specific operation?
|
| 138 |
+
// (If so, you'll want one of the FLOATING or INTEGRAL macros.)
|
| 139 |
+
//
|
| 140 |
+
// 2. Should half be supported? (If you're on CPU, the answer is almost
|
| 141 |
+
// definitely no. If you do want support, use one of the AND_HALF
|
| 142 |
+
// macros)
|
| 143 |
+
//
|
| 144 |
+
// Much rarer situations:
|
| 145 |
+
//
|
| 146 |
+
// 3. Should bool be supported? (You often have to write your kernel
|
| 147 |
+
// differently if arithmetic operations are involved.) If so,
|
| 148 |
+
// Use AT_DISPATCH_ALL_TYPES_AND along with ScalarType::Bool
|
| 149 |
+
//
|
| 150 |
+
// 4. Should complex be supported? The answer is almost always no,
|
| 151 |
+
// unless you are working on "generic" code that should work on
|
| 152 |
+
// all dtypes.
|
| 153 |
+
//
|
| 154 |
+
// Parameters:
|
| 155 |
+
// -----------
|
| 156 |
+
//
|
| 157 |
+
// 1. The NAME argument is a "tag" that is used to trace and then
|
| 158 |
+
// conditionally compile fragments of the case statements such
|
| 159 |
+
// that the kernel functions are specialized only for the dtypes
|
| 160 |
+
// that are needed. The NAME parameter *must* be a build time
|
| 161 |
+
// const char* (can't be std::string, etc...)
|
| 162 |
+
//
|
| 163 |
+
// Please ensure that the NAME is unique for every implementation
|
| 164 |
+
// or you run the risk of over-including code for the kernel
|
| 165 |
+
// functions. There is no risk of missing out on any code, so
|
| 166 |
+
// it's mostly a risk of a Type-2 error, and not a Type-1 error.
|
| 167 |
+
//
|
| 168 |
+
// Switch-like syntax:
|
| 169 |
+
// -------------------
|
| 170 |
+
// There is also a switch-case like syntax which is useful if a kernel
|
| 171 |
+
// needs to be specialized for particular scalar types
|
| 172 |
+
//
|
| 173 |
+
// AT_DISPATCH_SWITCH(self.scalar_type(), "op_name",
|
| 174 |
+
// AT_DISPATCH_CASE_INTEGRAL_TYPES([&] {
|
| 175 |
+
// op_integral<scalar_t>(iter);
|
| 176 |
+
// })
|
| 177 |
+
// AT_DISPATCH_CASE_FLOATING_TYPES([&] {
|
| 178 |
+
// op_floating<scalar_t>(iter);
|
| 179 |
+
// })
|
| 180 |
+
// AT_DISPATCH_CASE(kBool, [&] {
|
| 181 |
+
// op_bool(iter);
|
| 182 |
+
// })
|
| 183 |
+
// );
|
| 184 |
+
//
|
| 185 |
+
// For each AT_DISPATCH_FOO macro, there is a corresponding
|
| 186 |
+
// AT_DISPATCH_CASE_FOO macro which can be used inside of an
|
| 187 |
+
// AT_DISPATCH_SWITCH block.
|
| 188 |
+
|
| 189 |
+
// NB: the the_type variable is not used, but we have kept it for
|
| 190 |
+
// backwards compatibility. It's probably not used by anyone though;
|
| 191 |
+
// but we're just being safe (and it doesn't hurt.) Note we must
|
| 192 |
+
// use it to shut up warnings about unused store.
|
| 193 |
+
|
| 194 |
+
#define AT_DISPATCH_SWITCH(TYPE, NAME, ...) \
|
| 195 |
+
[&] { \
|
| 196 |
+
const auto& the_type = TYPE; \
|
| 197 |
+
constexpr const char* at_dispatch_name = NAME; \
|
| 198 |
+
/* don't use TYPE again in case it is an expensive or side-effect op */ \
|
| 199 |
+
at::ScalarType _st = ::detail::scalar_type(the_type); \
|
| 200 |
+
RECORD_KERNEL_FUNCTION_DTYPE(at_dispatch_name, _st); \
|
| 201 |
+
switch (_st) { \
|
| 202 |
+
__VA_ARGS__ \
|
| 203 |
+
default: \
|
| 204 |
+
TORCH_CHECK( \
|
| 205 |
+
false, \
|
| 206 |
+
'"', \
|
| 207 |
+
at_dispatch_name, \
|
| 208 |
+
"\" not implemented for '", \
|
| 209 |
+
toString(_st), \
|
| 210 |
+
"'"); \
|
| 211 |
+
} \
|
| 212 |
+
}()
|
| 213 |
+
|
| 214 |
+
#define AT_DISPATCH_CASE_FLOATING_TYPES(...) \
|
| 215 |
+
AT_DISPATCH_CASE(at::ScalarType::Double, __VA_ARGS__) \
|
| 216 |
+
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__)
|
| 217 |
+
|
| 218 |
+
#define AT_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
|
| 219 |
+
AT_DISPATCH_SWITCH(TYPE, NAME, AT_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
|
| 220 |
+
|
| 221 |
+
#define AT_DISPATCH_CASE_FLOATING_TYPES_AND_HALF(...) \
|
| 222 |
+
AT_DISPATCH_CASE(at::ScalarType::Double, __VA_ARGS__) \
|
| 223 |
+
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
|
| 224 |
+
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
|
| 225 |
+
|
| 226 |
+
#define AT_DISPATCH_FLOATING_TYPES_AND_HALF(TYPE, NAME, ...) \
|
| 227 |
+
AT_DISPATCH_SWITCH( \
|
| 228 |
+
TYPE, NAME, AT_DISPATCH_CASE_FLOATING_TYPES_AND_HALF(__VA_ARGS__))
|
| 229 |
+
|
| 230 |
+
#define AT_DISPATCH_CASE_REDUCED_FLOATING_TYPES(...) \
|
| 231 |
+
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
|
| 232 |
+
AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
|
| 233 |
+
|
| 234 |
+
#define AT_DISPATCH_REDUCED_FLOATING_TYPES(TYPE, NAME, ...) \
|
| 235 |
+
AT_DISPATCH_SWITCH( \
|
| 236 |
+
TYPE, NAME, AT_DISPATCH_CASE_REDUCED_FLOATING_TYPES(__VA_ARGS__))
|
| 237 |
+
|
| 238 |
+
#define AT_DISPATCH_CASE_FLOATING_TYPES_AND(SCALARTYPE, ...) \
|
| 239 |
+
AT_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__) \
|
| 240 |
+
AT_DISPATCH_CASE(SCALARTYPE, __VA_ARGS__)
|
| 241 |
+
|
| 242 |
+
#define AT_DISPATCH_FLOATING_TYPES_AND(SCALARTYPE, TYPE, NAME, ...) \
|
| 243 |
+
AT_DISPATCH_SWITCH( \
|
| 244 |
+
TYPE, \
|
| 245 |
+
NAME, \
|
| 246 |
+
AT_DISPATCH_CASE_FLOATING_TYPES_AND(SCALARTYPE, __VA_ARGS__))
|
| 247 |
+
|
| 248 |
+
#define AT_DISPATCH_CASE_FLOATING_TYPES_AND2(SCALARTYPE1, SCALARTYPE2, ...) \
|
| 249 |
+
AT_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__) \
|
| 250 |
+
AT_DISPATCH_CASE(SCALARTYPE1, __VA_ARGS__) \
|
| 251 |
+
AT_DISPATCH_CASE(SCALARTYPE2, __VA_ARGS__)
|
| 252 |
+
|
| 253 |
+
#define AT_DISPATCH_FLOATING_TYPES_AND2( \
|
| 254 |
+
SCALARTYPE1, SCALARTYPE2, TYPE, NAME, ...) \
|
| 255 |
+
AT_DISPATCH_SWITCH( \
|
| 256 |
+
TYPE, \
|
| 257 |
+
NAME, \
|
| 258 |
+
AT_DISPATCH_CASE_FLOATING_TYPES_AND2( \
|
| 259 |
+
SCALARTYPE1, SCALARTYPE2, __VA_ARGS__))
|
| 260 |
+
|
| 261 |
+
#define AT_DISPATCH_CASE_FLOATING_TYPES_AND3( \
|
| 262 |
+
SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, ...) \
|
| 263 |
+
AT_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__) \
|
| 264 |
+
AT_DISPATCH_CASE(SCALARTYPE1, __VA_ARGS__) \
|
| 265 |
+
AT_DISPATCH_CASE(SCALARTYPE2, __VA_ARGS__) \
|
| 266 |
+
AT_DISPATCH_CASE(SCALARTYPE3, __VA_ARGS__)
|
| 267 |
+
|
| 268 |
+
#define AT_DISPATCH_FLOATING_TYPES_AND3( \
|
| 269 |
+
SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, TYPE, NAME, ...) \
|
| 270 |
+
AT_DISPATCH_SWITCH( \
|
| 271 |
+
TYPE, \
|
| 272 |
+
NAME, \
|
| 273 |
+
AT_DISPATCH_CASE_FLOATING_TYPES_AND3( \
|
| 274 |
+
SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, __VA_ARGS__))
|
| 275 |
+
|
| 276 |
+
#define AT_DISPATCH_CASE_FLOATING_TYPES_AND4( \
|
| 277 |
+
SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, SCALARTYPE4, ...) \
|
| 278 |
+
AT_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__) \
|
| 279 |
+
AT_DISPATCH_CASE(SCALARTYPE1, __VA_ARGS__) \
|
| 280 |
+
AT_DISPATCH_CASE(SCALARTYPE2, __VA_ARGS__) \
|
| 281 |
+
AT_DISPATCH_CASE(SCALARTYPE3, __VA_ARGS__) \
|
| 282 |
+
AT_DISPATCH_CASE(SCALARTYPE4, __VA_ARGS__)
|
| 283 |
+
|
| 284 |
+
#define AT_DISPATCH_CASE_FLOATING_TYPES_AND5( \
|
| 285 |
+
SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, SCALARTYPE4, SCALARTYPE5, ...) \
|
| 286 |
+
AT_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__) \
|
| 287 |
+
AT_DISPATCH_CASE(SCALARTYPE1, __VA_ARGS__) \
|
| 288 |
+
AT_DISPATCH_CASE(SCALARTYPE2, __VA_ARGS__) \
|
| 289 |
+
AT_DISPATCH_CASE(SCALARTYPE3, __VA_ARGS__) \
|
| 290 |
+
AT_DISPATCH_CASE(SCALARTYPE4, __VA_ARGS__) \
|
| 291 |
+
AT_DISPATCH_CASE(SCALARTYPE5, __VA_ARGS__)
|
| 292 |
+
|
| 293 |
+
#define AT_DISPATCH_FLOATING_TYPES_AND4( \
|
| 294 |
+
SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, SCALARTYPE4, TYPE, NAME, ...) \
|
| 295 |
+
AT_DISPATCH_SWITCH( \
|
| 296 |
+
TYPE, \
|
| 297 |
+
NAME, \
|
| 298 |
+
AT_DISPATCH_CASE_FLOATING_TYPES_AND4( \
|
| 299 |
+
SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, SCALARTYPE4, __VA_ARGS__))
|
| 300 |
+
|
| 301 |
+
#define AT_DISPATCH_FLOATING_TYPES_AND5( \
|
| 302 |
+
SCALARTYPE1, \
|
| 303 |
+
SCALARTYPE2, \
|
| 304 |
+
SCALARTYPE3, \
|
| 305 |
+
SCALARTYPE4, \
|
| 306 |
+
SCALARTYPE5, \
|
| 307 |
+
TYPE, \
|
| 308 |
+
NAME, \
|
| 309 |
+
...) \
|
| 310 |
+
AT_DISPATCH_SWITCH( \
|
| 311 |
+
TYPE, \
|
| 312 |
+
NAME, \
|
| 313 |
+
AT_DISPATCH_CASE_FLOATING_TYPES_AND5( \
|
| 314 |
+
SCALARTYPE1, \
|
| 315 |
+
SCALARTYPE2, \
|
| 316 |
+
SCALARTYPE3, \
|
| 317 |
+
SCALARTYPE4, \
|
| 318 |
+
SCALARTYPE5, \
|
| 319 |
+
__VA_ARGS__))
|
| 320 |
+
|
| 321 |
+
#define AT_DISPATCH_CASE_COMPLEX_TYPES(...) \
|
| 322 |
+
AT_DISPATCH_CASE(at::ScalarType::ComplexDouble, __VA_ARGS__) \
|
| 323 |
+
AT_DISPATCH_CASE(at::ScalarType::ComplexFloat, __VA_ARGS__)
|
| 324 |
+
|
| 325 |
+
#define AT_DISPATCH_COMPLEX_TYPES(TYPE, NAME, ...) \
|
| 326 |
+
AT_DISPATCH_SWITCH(TYPE, NAME, AT_DISPATCH_CASE_COMPLEX_TYPES(__VA_ARGS__))
|
| 327 |
+
|
| 328 |
+
#define AT_DISPATCH_CASE_COMPLEX_TYPES_AND(SCALARTYPE, ...) \
|
| 329 |
+
AT_DISPATCH_CASE_COMPLEX_TYPES(__VA_ARGS__) \
|
| 330 |
+
AT_DISPATCH_CASE(SCALARTYPE, __VA_ARGS__)
|
| 331 |
+
|
| 332 |
+
#define AT_DISPATCH_COMPLEX_TYPES_AND(SCALARTYPE, TYPE, NAME, ...) \
|
| 333 |
+
AT_DISPATCH_SWITCH( \
|
| 334 |
+
TYPE, NAME, AT_DISPATCH_CASE_COMPLEX_TYPES_AND(SCALARTYPE, __VA_ARGS__))
|
| 335 |
+
|
| 336 |
+
#define AT_DISPATCH_CASE_FLOATING_AND_COMPLEX_TYPES(...) \
|
| 337 |
+
AT_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__) \
|
| 338 |
+
AT_DISPATCH_CASE_COMPLEX_TYPES(__VA_ARGS__)
|
| 339 |
+
|
| 340 |
+
#define AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(TYPE, NAME, ...) \
|
| 341 |
+
AT_DISPATCH_SWITCH( \
|
| 342 |
+
TYPE, NAME, AT_DISPATCH_CASE_FLOATING_AND_COMPLEX_TYPES(__VA_ARGS__))
|
| 343 |
+
|
| 344 |
+
#define AT_DISPATCH_CASE_FLOATING_AND_COMPLEX_TYPES_AND1(SCALARTYPE, ...) \
|
| 345 |
+
AT_DISPATCH_CASE_FLOATING_AND_COMPLEX_TYPES(__VA_ARGS__) \
|
| 346 |
+
AT_DISPATCH_CASE(SCALARTYPE, __VA_ARGS__)
|
| 347 |
+
|
| 348 |
+
#define AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND1( \
|
| 349 |
+
SCALARTYPE, TYPE, NAME, ...) \
|
| 350 |
+
AT_DISPATCH_SWITCH( \
|
| 351 |
+
TYPE, \
|
| 352 |
+
NAME, \
|
| 353 |
+
AT_DISPATCH_CASE_FLOATING_AND_COMPLEX_TYPES_AND1( \
|
| 354 |
+
SCALARTYPE, __VA_ARGS__))
|
| 355 |
+
|
| 356 |
+
#define AT_DISPATCH_CASE_FLOATING_AND_COMPLEX_TYPES_AND2( \
|
| 357 |
+
SCALARTYPE1, SCALARTYPE2, ...) \
|
| 358 |
+
AT_DISPATCH_CASE_FLOATING_AND_COMPLEX_TYPES(__VA_ARGS__) \
|
| 359 |
+
AT_DISPATCH_CASE(SCALARTYPE1, __VA_ARGS__) \
|
| 360 |
+
AT_DISPATCH_CASE(SCALARTYPE2, __VA_ARGS__)
|
| 361 |
+
|
| 362 |
+
#define AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND2( \
|
| 363 |
+
SCALARTYPE1, SCALARTYPE2, TYPE, NAME, ...) \
|
| 364 |
+
AT_DISPATCH_SWITCH( \
|
| 365 |
+
TYPE, \
|
| 366 |
+
NAME, \
|
| 367 |
+
AT_DISPATCH_CASE_FLOATING_AND_COMPLEX_TYPES_AND2( \
|
| 368 |
+
SCALARTYPE1, SCALARTYPE2, __VA_ARGS__))
|
| 369 |
+
|
| 370 |
+
#define AT_DISPATCH_CASE_FLOATING_AND_COMPLEX_TYPES_AND3( \
|
| 371 |
+
SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, ...) \
|
| 372 |
+
AT_DISPATCH_CASE_FLOATING_AND_COMPLEX_TYPES(__VA_ARGS__) \
|
| 373 |
+
AT_DISPATCH_CASE(SCALARTYPE1, __VA_ARGS__) \
|
| 374 |
+
AT_DISPATCH_CASE(SCALARTYPE2, __VA_ARGS__) \
|
| 375 |
+
AT_DISPATCH_CASE(SCALARTYPE3, __VA_ARGS__)
|
| 376 |
+
|
| 377 |
+
#define AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND3( \
|
| 378 |
+
SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, TYPE, NAME, ...) \
|
| 379 |
+
AT_DISPATCH_SWITCH( \
|
| 380 |
+
TYPE, \
|
| 381 |
+
NAME, \
|
| 382 |
+
AT_DISPATCH_CASE_FLOATING_AND_COMPLEX_TYPES_AND3( \
|
| 383 |
+
SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, __VA_ARGS__))
|
| 384 |
+
|
| 385 |
+
#define AT_DISPATCH_CASE_FLOATING_AND_COMPLEX_TYPES_AND4( \
|
| 386 |
+
SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, SCALARTYPE4, ...) \
|
| 387 |
+
AT_DISPATCH_CASE_FLOATING_AND_COMPLEX_TYPES(__VA_ARGS__) \
|
| 388 |
+
AT_DISPATCH_CASE(SCALARTYPE1, __VA_ARGS__) \
|
| 389 |
+
AT_DISPATCH_CASE(SCALARTYPE2, __VA_ARGS__) \
|
| 390 |
+
AT_DISPATCH_CASE(SCALARTYPE3, __VA_ARGS__) \
|
| 391 |
+
AT_DISPATCH_CASE(SCALARTYPE4, __VA_ARGS__)
|
| 392 |
+
|
| 393 |
+
#define AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND4( \
|
| 394 |
+
SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, SCALARTYPE4, TYPE, NAME, ...) \
|
| 395 |
+
AT_DISPATCH_SWITCH( \
|
| 396 |
+
TYPE, \
|
| 397 |
+
NAME, \
|
| 398 |
+
AT_DISPATCH_CASE_FLOATING_AND_COMPLEX_TYPES_AND4( \
|
| 399 |
+
SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, SCALARTYPE4, __VA_ARGS__))
|
| 400 |
+
|
| 401 |
+
#define AT_DISPATCH_CASE_FLOATING_AND_COMPLEX_TYPES_AND5( \
|
| 402 |
+
SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, SCALARTYPE4, SCALARTYPE5, ...) \
|
| 403 |
+
AT_DISPATCH_CASE_FLOATING_AND_COMPLEX_TYPES(__VA_ARGS__) \
|
| 404 |
+
AT_DISPATCH_CASE(SCALARTYPE1, __VA_ARGS__) \
|
| 405 |
+
AT_DISPATCH_CASE(SCALARTYPE2, __VA_ARGS__) \
|
| 406 |
+
AT_DISPATCH_CASE(SCALARTYPE3, __VA_ARGS__) \
|
| 407 |
+
AT_DISPATCH_CASE(SCALARTYPE4, __VA_ARGS__) \
|
| 408 |
+
AT_DISPATCH_CASE(SCALARTYPE5, __VA_ARGS__)
|
| 409 |
+
|
| 410 |
+
#define AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND5( \
|
| 411 |
+
SCALARTYPE1, \
|
| 412 |
+
SCALARTYPE2, \
|
| 413 |
+
SCALARTYPE3, \
|
| 414 |
+
SCALARTYPE4, \
|
| 415 |
+
SCALARTYPE5, \
|
| 416 |
+
TYPE, \
|
| 417 |
+
NAME, \
|
| 418 |
+
...) \
|
| 419 |
+
AT_DISPATCH_SWITCH( \
|
| 420 |
+
TYPE, \
|
| 421 |
+
NAME, \
|
| 422 |
+
AT_DISPATCH_CASE_FLOATING_AND_COMPLEX_TYPES_AND5( \
|
| 423 |
+
SCALARTYPE1, \
|
| 424 |
+
SCALARTYPE2, \
|
| 425 |
+
SCALARTYPE3, \
|
| 426 |
+
SCALARTYPE4, \
|
| 427 |
+
SCALARTYPE5, \
|
| 428 |
+
__VA_ARGS__))
|
| 429 |
+
|
| 430 |
+
#define AT_DISPATCH_CASE_FLOATING_AND_COMPLEX_TYPES_AND6( \
|
| 431 |
+
SCALARTYPE1, \
|
| 432 |
+
SCALARTYPE2, \
|
| 433 |
+
SCALARTYPE3, \
|
| 434 |
+
SCALARTYPE4, \
|
| 435 |
+
SCALARTYPE5, \
|
| 436 |
+
SCALARTYPE6, \
|
| 437 |
+
...) \
|
| 438 |
+
AT_DISPATCH_CASE_FLOATING_AND_COMPLEX_TYPES(__VA_ARGS__) \
|
| 439 |
+
AT_DISPATCH_CASE(SCALARTYPE1, __VA_ARGS__) \
|
| 440 |
+
AT_DISPATCH_CASE(SCALARTYPE2, __VA_ARGS__) \
|
| 441 |
+
AT_DISPATCH_CASE(SCALARTYPE3, __VA_ARGS__) \
|
| 442 |
+
AT_DISPATCH_CASE(SCALARTYPE4, __VA_ARGS__) \
|
| 443 |
+
AT_DISPATCH_CASE(SCALARTYPE5, __VA_ARGS__) \
|
| 444 |
+
AT_DISPATCH_CASE(SCALARTYPE6, __VA_ARGS__)
|
| 445 |
+
|
| 446 |
+
#define AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND6( \
|
| 447 |
+
SCALARTYPE1, \
|
| 448 |
+
SCALARTYPE2, \
|
| 449 |
+
SCALARTYPE3, \
|
| 450 |
+
SCALARTYPE4, \
|
| 451 |
+
SCALARTYPE5, \
|
| 452 |
+
SCALARTYPE6, \
|
| 453 |
+
TYPE, \
|
| 454 |
+
NAME, \
|
| 455 |
+
...) \
|
| 456 |
+
AT_DISPATCH_SWITCH( \
|
| 457 |
+
TYPE, \
|
| 458 |
+
NAME, \
|
| 459 |
+
AT_DISPATCH_CASE_FLOATING_AND_COMPLEX_TYPES_AND6( \
|
| 460 |
+
SCALARTYPE1, \
|
| 461 |
+
SCALARTYPE2, \
|
| 462 |
+
SCALARTYPE3, \
|
| 463 |
+
SCALARTYPE4, \
|
| 464 |
+
SCALARTYPE5, \
|
| 465 |
+
SCALARTYPE6, \
|
| 466 |
+
__VA_ARGS__))
|
| 467 |
+
|
| 468 |
+
#define AT_DISPATCH_CASE_INTEGRAL_TYPES(...) \
|
| 469 |
+
AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__) \
|
| 470 |
+
AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__) \
|
| 471 |
+
AT_DISPATCH_CASE(at::ScalarType::Int, __VA_ARGS__) \
|
| 472 |
+
AT_DISPATCH_CASE(at::ScalarType::Long, __VA_ARGS__) \
|
| 473 |
+
AT_DISPATCH_CASE(at::ScalarType::Short, __VA_ARGS__)
|
| 474 |
+
|
| 475 |
+
#define AT_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...) \
|
| 476 |
+
AT_DISPATCH_SWITCH(TYPE, NAME, AT_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__))
|
| 477 |
+
|
| 478 |
+
#define AT_DISPATCH_CASE_INTEGRAL_TYPES_AND(SCALARTYPE, ...) \
|
| 479 |
+
AT_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__) \
|
| 480 |
+
AT_DISPATCH_CASE(SCALARTYPE, __VA_ARGS__)
|
| 481 |
+
|
| 482 |
+
#define AT_DISPATCH_INTEGRAL_TYPES_AND(SCALARTYPE, TYPE, NAME, ...) \
|
| 483 |
+
AT_DISPATCH_SWITCH( \
|
| 484 |
+
TYPE, \
|
| 485 |
+
NAME, \
|
| 486 |
+
AT_DISPATCH_CASE_INTEGRAL_TYPES_AND(SCALARTYPE, __VA_ARGS__))
|
| 487 |
+
|
| 488 |
+
#define AT_DISPATCH_CASE_ALL_TYPES(...) \
|
| 489 |
+
AT_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__) \
|
| 490 |
+
AT_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)
|
| 491 |
+
|
| 492 |
+
#define AT_DISPATCH_ALL_TYPES(TYPE, NAME, ...) \
|
| 493 |
+
AT_DISPATCH_SWITCH(TYPE, NAME, AT_DISPATCH_CASE_ALL_TYPES(__VA_ARGS__))
|
| 494 |
+
|
| 495 |
+
#define AT_DISPATCH_CASE_QINT_TYPES(...) \
|
| 496 |
+
AT_DISPATCH_CASE_QINT(at::kQInt8, at::qint8, __VA_ARGS__) \
|
| 497 |
+
AT_DISPATCH_CASE_QINT(at::kQUInt8, at::quint8, __VA_ARGS__) \
|
| 498 |
+
AT_DISPATCH_CASE_QINT(at::kQInt32, at::qint32, __VA_ARGS__)
|
| 499 |
+
|
| 500 |
+
#define AT_DISPATCH_QINT_TYPES(TYPE, NAME, ...) \
|
| 501 |
+
AT_DISPATCH_SWITCH(TYPE, NAME, AT_DISPATCH_CASE_QINT_TYPES(__VA_ARGS__))
|
| 502 |
+
|
| 503 |
+
#define AT_DISPATCH_CASE_QINT_TYPES_AND(SCALARTYPE, ...) \
|
| 504 |
+
AT_DISPATCH_CASE_QINT_TYPES(__VA_ARGS__) \
|
| 505 |
+
AT_DISPATCH_CASE(SCALARTYPE, __VA_ARGS__)
|
| 506 |
+
|
| 507 |
+
#define AT_DISPATCH_QINT_TYPES_AND(SCALARTYPE, TYPE, NAME, ...) \
|
| 508 |
+
AT_DISPATCH_SWITCH( \
|
| 509 |
+
TYPE, NAME, AT_DISPATCH_CASE_QINT_TYPES_AND(SCALARTYPE, __VA_ARGS__))
|
| 510 |
+
|
| 511 |
+
#define AT_DISPATCH_CASE_QINT_BYTE_TYPES(...) \
|
| 512 |
+
AT_DISPATCH_CASE_QINT(at::kQInt8, at::qint8, __VA_ARGS__) \
|
| 513 |
+
AT_DISPATCH_CASE_QINT(at::kQUInt8, at::quint8, __VA_ARGS__)
|
| 514 |
+
|
| 515 |
+
#define AT_DISPATCH_QINT_BYTE_TYPES(TYPE, NAME, ...) \
|
| 516 |
+
AT_DISPATCH_SWITCH(TYPE, NAME, AT_DISPATCH_CASE_QINT_BYTE_TYPES(__VA_ARGS__))
|
| 517 |
+
|
| 518 |
+
#define AT_DISPATCH_CASE_QINT_AND_SUB_BYTE_TYPES(...) \
|
| 519 |
+
AT_QINT_SUB_BYTE_PRIVATE_CASE_TYPE( \
|
| 520 |
+
at::kQInt8, at::qint8, CHAR_BIT, SCHAR_MIN, SCHAR_MAX, __VA_ARGS__) \
|
| 521 |
+
AT_QINT_SUB_BYTE_PRIVATE_CASE_TYPE( \
|
| 522 |
+
at::kQUInt8, at::quint8, CHAR_BIT, 0, UCHAR_MAX, __VA_ARGS__) \
|
| 523 |
+
AT_QINT_SUB_BYTE_PRIVATE_CASE_TYPE( \
|
| 524 |
+
at::kQInt32, \
|
| 525 |
+
at::qint32, \
|
| 526 |
+
CHAR_BIT * sizeof(int), \
|
| 527 |
+
INT_MIN, \
|
| 528 |
+
INT_MAX, \
|
| 529 |
+
__VA_ARGS__) \
|
| 530 |
+
AT_QINT_SUB_BYTE_PRIVATE_CASE_TYPE( \
|
| 531 |
+
at::kQUInt4x2, at::quint4x2, 4, 0, 15, __VA_ARGS__) \
|
| 532 |
+
AT_QINT_SUB_BYTE_PRIVATE_CASE_TYPE( \
|
| 533 |
+
at::kQUInt2x4, at::quint2x4, 2, 0, 3, __VA_ARGS__)
|
| 534 |
+
|
| 535 |
+
#define AT_DISPATCH_QINT_AND_SUB_BYTE_TYPES(TYPE, NAME, ...) \
|
| 536 |
+
AT_DISPATCH_SWITCH( \
|
| 537 |
+
TYPE, NAME, AT_DISPATCH_CASE_QINT_AND_SUB_BYTE_TYPES(__VA_ARGS__))
|
| 538 |
+
|
| 539 |
+
#define AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX(...) \
|
| 540 |
+
AT_DISPATCH_CASE_ALL_TYPES(__VA_ARGS__) \
|
| 541 |
+
AT_DISPATCH_CASE_COMPLEX_TYPES(__VA_ARGS__)
|
| 542 |
+
|
| 543 |
+
#define AT_DISPATCH_ALL_TYPES_AND_COMPLEX(TYPE, NAME, ...) \
|
| 544 |
+
AT_DISPATCH_SWITCH( \
|
| 545 |
+
TYPE, NAME, AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX(__VA_ARGS__))
|
| 546 |
+
|
| 547 |
+
#define AT_DISPATCH_CASE_ALL_TYPES_AND(SCALARTYPE, ...) \
|
| 548 |
+
AT_DISPATCH_CASE_ALL_TYPES(__VA_ARGS__) \
|
| 549 |
+
AT_DISPATCH_CASE(SCALARTYPE, __VA_ARGS__)
|
| 550 |
+
|
| 551 |
+
#define AT_DISPATCH_ALL_TYPES_AND(SCALARTYPE, TYPE, NAME, ...) \
|
| 552 |
+
AT_DISPATCH_SWITCH( \
|
| 553 |
+
TYPE, NAME, AT_DISPATCH_CASE_ALL_TYPES_AND(SCALARTYPE, __VA_ARGS__))
|
| 554 |
+
|
| 555 |
+
#define AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX_AND(SCALARTYPE, ...) \
|
| 556 |
+
AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX(__VA_ARGS__) \
|
| 557 |
+
AT_DISPATCH_CASE(SCALARTYPE, __VA_ARGS__)
|
| 558 |
+
|
| 559 |
+
#define AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND(SCALARTYPE, TYPE, NAME, ...) \
|
| 560 |
+
AT_DISPATCH_SWITCH( \
|
| 561 |
+
TYPE, \
|
| 562 |
+
NAME, \
|
| 563 |
+
AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX_AND(SCALARTYPE, __VA_ARGS__))
|
| 564 |
+
|
| 565 |
+
#define AT_DISPATCH_CASE_ALL_TYPES_AND2(SCALARTYPE1, SCALARTYPE2, ...) \
|
| 566 |
+
AT_DISPATCH_CASE_ALL_TYPES(__VA_ARGS__) \
|
| 567 |
+
AT_DISPATCH_CASE(SCALARTYPE1, __VA_ARGS__) \
|
| 568 |
+
AT_DISPATCH_CASE(SCALARTYPE2, __VA_ARGS__)
|
| 569 |
+
|
| 570 |
+
#define AT_DISPATCH_ALL_TYPES_AND2(SCALARTYPE1, SCALARTYPE2, TYPE, NAME, ...) \
|
| 571 |
+
AT_DISPATCH_SWITCH( \
|
| 572 |
+
TYPE, \
|
| 573 |
+
NAME, \
|
| 574 |
+
AT_DISPATCH_CASE_ALL_TYPES_AND2(SCALARTYPE1, SCALARTYPE2, __VA_ARGS__))
|
| 575 |
+
|
| 576 |
+
#define AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX_AND2( \
|
| 577 |
+
SCALARTYPE1, SCALARTYPE2, ...) \
|
| 578 |
+
AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX(__VA_ARGS__) \
|
| 579 |
+
AT_DISPATCH_CASE(SCALARTYPE1, __VA_ARGS__) \
|
| 580 |
+
AT_DISPATCH_CASE(SCALARTYPE2, __VA_ARGS__)
|
| 581 |
+
|
| 582 |
+
#define AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2( \
|
| 583 |
+
SCALARTYPE1, SCALARTYPE2, TYPE, NAME, ...) \
|
| 584 |
+
AT_DISPATCH_SWITCH( \
|
| 585 |
+
TYPE, \
|
| 586 |
+
NAME, \
|
| 587 |
+
AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX_AND2( \
|
| 588 |
+
SCALARTYPE1, SCALARTYPE2, __VA_ARGS__))
|
| 589 |
+
|
| 590 |
+
#define AT_DISPATCH_CASE_ALL_TYPES_AND3( \
|
| 591 |
+
SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, ...) \
|
| 592 |
+
AT_DISPATCH_CASE_ALL_TYPES(__VA_ARGS__) \
|
| 593 |
+
AT_DISPATCH_CASE(SCALARTYPE1, __VA_ARGS__) \
|
| 594 |
+
AT_DISPATCH_CASE(SCALARTYPE2, __VA_ARGS__) \
|
| 595 |
+
AT_DISPATCH_CASE(SCALARTYPE3, __VA_ARGS__)
|
| 596 |
+
|
| 597 |
+
#define AT_DISPATCH_ALL_TYPES_AND3( \
|
| 598 |
+
SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, TYPE, NAME, ...) \
|
| 599 |
+
AT_DISPATCH_SWITCH( \
|
| 600 |
+
TYPE, \
|
| 601 |
+
NAME, \
|
| 602 |
+
AT_DISPATCH_CASE_ALL_TYPES_AND3( \
|
| 603 |
+
SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, __VA_ARGS__))
|
| 604 |
+
|
| 605 |
+
#define AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX_AND3( \
|
| 606 |
+
SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, ...) \
|
| 607 |
+
AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX(__VA_ARGS__) \
|
| 608 |
+
AT_DISPATCH_CASE(SCALARTYPE1, __VA_ARGS__) \
|
| 609 |
+
AT_DISPATCH_CASE(SCALARTYPE2, __VA_ARGS__) \
|
| 610 |
+
AT_DISPATCH_CASE(SCALARTYPE3, __VA_ARGS__)
|
| 611 |
+
|
| 612 |
+
#define AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND3( \
|
| 613 |
+
SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, TYPE, NAME, ...) \
|
| 614 |
+
AT_DISPATCH_SWITCH( \
|
| 615 |
+
TYPE, \
|
| 616 |
+
NAME, \
|
| 617 |
+
AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX_AND3( \
|
| 618 |
+
SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, __VA_ARGS__))
|
| 619 |
+
|
| 620 |
+
#define AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX_AND4( \
|
| 621 |
+
SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, SCALARTYPE4, ...) \
|
| 622 |
+
AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX(__VA_ARGS__) \
|
| 623 |
+
AT_DISPATCH_CASE(SCALARTYPE1, __VA_ARGS__) \
|
| 624 |
+
AT_DISPATCH_CASE(SCALARTYPE2, __VA_ARGS__) \
|
| 625 |
+
AT_DISPATCH_CASE(SCALARTYPE3, __VA_ARGS__) \
|
| 626 |
+
AT_DISPATCH_CASE(SCALARTYPE4, __VA_ARGS__)
|
| 627 |
+
|
| 628 |
+
#define AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND4( \
|
| 629 |
+
SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, SCALARTYPE4, TYPE, NAME, ...) \
|
| 630 |
+
AT_DISPATCH_SWITCH( \
|
| 631 |
+
TYPE, \
|
| 632 |
+
NAME, \
|
| 633 |
+
AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX_AND4( \
|
| 634 |
+
SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, SCALARTYPE4, __VA_ARGS__))
|
| 635 |
+
|
| 636 |
+
#define AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX_AND5( \
|
| 637 |
+
SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, SCALARTYPE4, SCALARTYPE5, ...) \
|
| 638 |
+
AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX(__VA_ARGS__) \
|
| 639 |
+
AT_DISPATCH_CASE(SCALARTYPE1, __VA_ARGS__) \
|
| 640 |
+
AT_DISPATCH_CASE(SCALARTYPE2, __VA_ARGS__) \
|
| 641 |
+
AT_DISPATCH_CASE(SCALARTYPE3, __VA_ARGS__) \
|
| 642 |
+
AT_DISPATCH_CASE(SCALARTYPE4, __VA_ARGS__) \
|
| 643 |
+
AT_DISPATCH_CASE(SCALARTYPE5, __VA_ARGS__)
|
| 644 |
+
|
| 645 |
+
#define AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND5( \
|
| 646 |
+
SCALARTYPE1, \
|
| 647 |
+
SCALARTYPE2, \
|
| 648 |
+
SCALARTYPE3, \
|
| 649 |
+
SCALARTYPE4, \
|
| 650 |
+
SCALARTYPE5, \
|
| 651 |
+
TYPE, \
|
| 652 |
+
NAME, \
|
| 653 |
+
...) \
|
| 654 |
+
AT_DISPATCH_SWITCH( \
|
| 655 |
+
TYPE, \
|
| 656 |
+
NAME, \
|
| 657 |
+
AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX_AND5( \
|
| 658 |
+
SCALARTYPE1, \
|
| 659 |
+
SCALARTYPE2, \
|
| 660 |
+
SCALARTYPE3, \
|
| 661 |
+
SCALARTYPE4, \
|
| 662 |
+
SCALARTYPE5, \
|
| 663 |
+
__VA_ARGS__))
|
| 664 |
+
|
| 665 |
+
#define AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX_AND6( \
|
| 666 |
+
SCALARTYPE1, \
|
| 667 |
+
SCALARTYPE2, \
|
| 668 |
+
SCALARTYPE3, \
|
| 669 |
+
SCALARTYPE4, \
|
| 670 |
+
SCALARTYPE5, \
|
| 671 |
+
SCALARTYPE6, \
|
| 672 |
+
...) \
|
| 673 |
+
AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX(__VA_ARGS__) \
|
| 674 |
+
AT_DISPATCH_CASE(SCALARTYPE1, __VA_ARGS__) \
|
| 675 |
+
AT_DISPATCH_CASE(SCALARTYPE2, __VA_ARGS__) \
|
| 676 |
+
AT_DISPATCH_CASE(SCALARTYPE3, __VA_ARGS__) \
|
| 677 |
+
AT_DISPATCH_CASE(SCALARTYPE4, __VA_ARGS__) \
|
| 678 |
+
AT_DISPATCH_CASE(SCALARTYPE5, __VA_ARGS__) \
|
| 679 |
+
AT_DISPATCH_CASE(SCALARTYPE6, __VA_ARGS__)
|
| 680 |
+
|
| 681 |
+
#define AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND6( \
|
| 682 |
+
SCALARTYPE1, \
|
| 683 |
+
SCALARTYPE2, \
|
| 684 |
+
SCALARTYPE3, \
|
| 685 |
+
SCALARTYPE4, \
|
| 686 |
+
SCALARTYPE5, \
|
| 687 |
+
SCALARTYPE6, \
|
| 688 |
+
TYPE, \
|
| 689 |
+
NAME, \
|
| 690 |
+
...) \
|
| 691 |
+
AT_DISPATCH_SWITCH( \
|
| 692 |
+
TYPE, \
|
| 693 |
+
NAME, \
|
| 694 |
+
AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX_AND6( \
|
| 695 |
+
SCALARTYPE1, \
|
| 696 |
+
SCALARTYPE2, \
|
| 697 |
+
SCALARTYPE3, \
|
| 698 |
+
SCALARTYPE4, \
|
| 699 |
+
SCALARTYPE5, \
|
| 700 |
+
SCALARTYPE6, \
|
| 701 |
+
__VA_ARGS__))
|
| 702 |
+
|
| 703 |
+
#define AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX_AND7( \
|
| 704 |
+
SCALARTYPE1, \
|
| 705 |
+
SCALARTYPE2, \
|
| 706 |
+
SCALARTYPE3, \
|
| 707 |
+
SCALARTYPE4, \
|
| 708 |
+
SCALARTYPE5, \
|
| 709 |
+
SCALARTYPE6, \
|
| 710 |
+
SCALARTYPE7, \
|
| 711 |
+
...) \
|
| 712 |
+
AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX(__VA_ARGS__) \
|
| 713 |
+
AT_DISPATCH_CASE(SCALARTYPE1, __VA_ARGS__) \
|
| 714 |
+
AT_DISPATCH_CASE(SCALARTYPE2, __VA_ARGS__) \
|
| 715 |
+
AT_DISPATCH_CASE(SCALARTYPE3, __VA_ARGS__) \
|
| 716 |
+
AT_DISPATCH_CASE(SCALARTYPE4, __VA_ARGS__) \
|
| 717 |
+
AT_DISPATCH_CASE(SCALARTYPE5, __VA_ARGS__) \
|
| 718 |
+
AT_DISPATCH_CASE(SCALARTYPE6, __VA_ARGS__) \
|
| 719 |
+
AT_DISPATCH_CASE(SCALARTYPE7, __VA_ARGS__)
|
| 720 |
+
|
| 721 |
+
#define AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND7( \
|
| 722 |
+
SCALARTYPE1, \
|
| 723 |
+
SCALARTYPE2, \
|
| 724 |
+
SCALARTYPE3, \
|
| 725 |
+
SCALARTYPE4, \
|
| 726 |
+
SCALARTYPE5, \
|
| 727 |
+
SCALARTYPE6, \
|
| 728 |
+
SCALARTYPE7, \
|
| 729 |
+
TYPE, \
|
| 730 |
+
NAME, \
|
| 731 |
+
...) \
|
| 732 |
+
AT_DISPATCH_SWITCH( \
|
| 733 |
+
TYPE, \
|
| 734 |
+
NAME, \
|
| 735 |
+
AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX_AND7( \
|
| 736 |
+
SCALARTYPE1, \
|
| 737 |
+
SCALARTYPE2, \
|
| 738 |
+
SCALARTYPE3, \
|
| 739 |
+
SCALARTYPE4, \
|
| 740 |
+
SCALARTYPE5, \
|
| 741 |
+
SCALARTYPE6, \
|
| 742 |
+
SCALARTYPE7, \
|
| 743 |
+
__VA_ARGS__))
|
| 744 |
+
|
| 745 |
+
#define AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX_AND8( \
|
| 746 |
+
SCALARTYPE1, \
|
| 747 |
+
SCALARTYPE2, \
|
| 748 |
+
SCALARTYPE3, \
|
| 749 |
+
SCALARTYPE4, \
|
| 750 |
+
SCALARTYPE5, \
|
| 751 |
+
SCALARTYPE6, \
|
| 752 |
+
SCALARTYPE7, \
|
| 753 |
+
SCALARTYPE8, \
|
| 754 |
+
...) \
|
| 755 |
+
AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX(__VA_ARGS__) \
|
| 756 |
+
AT_DISPATCH_CASE(SCALARTYPE1, __VA_ARGS__) \
|
| 757 |
+
AT_DISPATCH_CASE(SCALARTYPE2, __VA_ARGS__) \
|
| 758 |
+
AT_DISPATCH_CASE(SCALARTYPE3, __VA_ARGS__) \
|
| 759 |
+
AT_DISPATCH_CASE(SCALARTYPE4, __VA_ARGS__) \
|
| 760 |
+
AT_DISPATCH_CASE(SCALARTYPE5, __VA_ARGS__) \
|
| 761 |
+
AT_DISPATCH_CASE(SCALARTYPE6, __VA_ARGS__) \
|
| 762 |
+
AT_DISPATCH_CASE(SCALARTYPE7, __VA_ARGS__) \
|
| 763 |
+
AT_DISPATCH_CASE(SCALARTYPE8, __VA_ARGS__)
|
| 764 |
+
|
| 765 |
+
#define AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND8( \
|
| 766 |
+
SCALARTYPE1, \
|
| 767 |
+
SCALARTYPE2, \
|
| 768 |
+
SCALARTYPE3, \
|
| 769 |
+
SCALARTYPE4, \
|
| 770 |
+
SCALARTYPE5, \
|
| 771 |
+
SCALARTYPE6, \
|
| 772 |
+
SCALARTYPE7, \
|
| 773 |
+
SCALARTYPE8, \
|
| 774 |
+
TYPE, \
|
| 775 |
+
NAME, \
|
| 776 |
+
...) \
|
| 777 |
+
AT_DISPATCH_SWITCH( \
|
| 778 |
+
TYPE, \
|
| 779 |
+
NAME, \
|
| 780 |
+
AT_DISPATCH_CASE_ALL_TYPES_AND_COMPLEX_AND8( \
|
| 781 |
+
SCALARTYPE1, \
|
| 782 |
+
SCALARTYPE2, \
|
| 783 |
+
SCALARTYPE3, \
|
| 784 |
+
SCALARTYPE4, \
|
| 785 |
+
SCALARTYPE5, \
|
| 786 |
+
SCALARTYPE6, \
|
| 787 |
+
SCALARTYPE7, \
|
| 788 |
+
SCALARTYPE8, \
|
| 789 |
+
__VA_ARGS__))
|
| 790 |
+
|
| 791 |
+
#define AT_DISPATCH_CASE_BIT_TYPES(...) \
|
| 792 |
+
AT_DISPATCH_CASE(at::ScalarType::Bits1x8, __VA_ARGS__) \
|
| 793 |
+
AT_DISPATCH_CASE(at::ScalarType::Bits2x4, __VA_ARGS__) \
|
| 794 |
+
AT_DISPATCH_CASE(at::ScalarType::Bits4x2, __VA_ARGS__) \
|
| 795 |
+
AT_DISPATCH_CASE(at::ScalarType::Bits8, __VA_ARGS__) \
|
| 796 |
+
AT_DISPATCH_CASE(at::ScalarType::Bits16, __VA_ARGS__)
|
| 797 |
+
|
| 798 |
+
#define AT_DISPATCH_BIT_TYPES(TYPE, NAME, ...) \
|
| 799 |
+
AT_DISPATCH_SWITCH(TYPE, NAME, AT_DISPATCH_CASE_BIT_TYPES(__VA_ARGS__))
|
| 800 |
+
|
| 801 |
+
#define AT_DISPATCH_INDEX_TYPES(TYPE, NAME, ...) \
|
| 802 |
+
AT_DISPATCH_SWITCH( \
|
| 803 |
+
TYPE, \
|
| 804 |
+
NAME, \
|
| 805 |
+
AT_PRIVATE_CASE_TYPE_USING_HINT( \
|
| 806 |
+
at::ScalarType::Int, index_t, __VA_ARGS__) \
|
| 807 |
+
AT_PRIVATE_CASE_TYPE_USING_HINT( \
|
| 808 |
+
at::ScalarType::Long, index_t, __VA_ARGS__))
|
lib/python3.12/site-packages/torch/include/ATen/ExpandBase.h
ADDED
|
@@ -0,0 +1,30 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#include <ATen/core/TensorBase.h>
|
| 2 |
+
|
| 3 |
+
// Broadcasting utilities for working with TensorBase
|
| 4 |
+
namespace at {
|
| 5 |
+
namespace internal {
|
| 6 |
+
TORCH_API TensorBase expand_slow_path(const TensorBase& self, IntArrayRef size);
|
| 7 |
+
} // namespace internal
|
| 8 |
+
|
| 9 |
+
inline c10::MaybeOwned<TensorBase> expand_size(
|
| 10 |
+
const TensorBase& self,
|
| 11 |
+
IntArrayRef size) {
|
| 12 |
+
if (size.equals(self.sizes())) {
|
| 13 |
+
return c10::MaybeOwned<TensorBase>::borrowed(self);
|
| 14 |
+
}
|
| 15 |
+
return c10::MaybeOwned<TensorBase>::owned(
|
| 16 |
+
at::internal::expand_slow_path(self, size));
|
| 17 |
+
}
|
| 18 |
+
c10::MaybeOwned<TensorBase> expand_size(TensorBase&& self, IntArrayRef size) =
|
| 19 |
+
delete;
|
| 20 |
+
|
| 21 |
+
inline c10::MaybeOwned<TensorBase> expand_inplace(
|
| 22 |
+
const TensorBase& tensor,
|
| 23 |
+
const TensorBase& to_expand) {
|
| 24 |
+
return expand_size(to_expand, tensor.sizes());
|
| 25 |
+
}
|
| 26 |
+
c10::MaybeOwned<TensorBase> expand_inplace(
|
| 27 |
+
const TensorBase& tensor,
|
| 28 |
+
TensorBase&& to_expand) = delete;
|
| 29 |
+
|
| 30 |
+
} // namespace at
|
lib/python3.12/site-packages/torch/include/ATen/Generator.h
ADDED
|
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
#include <ATen/core/Generator.h>
|
lib/python3.12/site-packages/torch/include/ATen/LegacyBatchedFallback.h
ADDED
|
@@ -0,0 +1,25 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
#include <ATen/ATen.h>
|
| 3 |
+
#include <ATen/core/op_registration/op_registration.h>
|
| 4 |
+
#include <torch/library.h>
|
| 5 |
+
|
| 6 |
+
namespace at {
|
| 7 |
+
|
| 8 |
+
// If an operator doesn't have a batching rule implemented then we fallback
|
| 9 |
+
// to this implementation. The fallback only works on out-of-place operators
|
| 10 |
+
// that return only tensors with new memory. (e.g., no in-place operators, no
|
| 11 |
+
// view operations).
|
| 12 |
+
//
|
| 13 |
+
// The fallback effectively takes all of the BatchedTensors in `stack`, slices
|
| 14 |
+
// them, and runs `op` on all of the corresponding slices to produce slices
|
| 15 |
+
// of the outputs. The output slices then get `torch.stack`ed to create the
|
| 16 |
+
// final returns.
|
| 17 |
+
//
|
| 18 |
+
// The performance of the fallback is not very good because it introduces an
|
| 19 |
+
// extra copy from stacking the sliced outputs. Because of this, we prefer to
|
| 20 |
+
// write batching rules for operators whenever possible.
|
| 21 |
+
void batchedTensorForLoopFallback(
|
| 22 |
+
const c10::OperatorHandle& op,
|
| 23 |
+
torch::jit::Stack* stack);
|
| 24 |
+
|
| 25 |
+
} // namespace at
|
lib/python3.12/site-packages/torch/include/ATen/MatrixRef.h
ADDED
|
@@ -0,0 +1,108 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
#include <ATen/Utils.h>
|
| 3 |
+
#include <c10/util/ArrayRef.h>
|
| 4 |
+
|
| 5 |
+
namespace at {
|
| 6 |
+
/// MatrixRef - Like an ArrayRef, but with an extra recorded strides so that
|
| 7 |
+
/// we can easily view it as a multidimensional array.
|
| 8 |
+
///
|
| 9 |
+
/// Like ArrayRef, this class does not own the underlying data, it is expected
|
| 10 |
+
/// to be used in situations where the data resides in some other buffer.
|
| 11 |
+
///
|
| 12 |
+
/// This is intended to be trivially copyable, so it should be passed by
|
| 13 |
+
/// value.
|
| 14 |
+
///
|
| 15 |
+
/// For now, 2D only (so the copies are actually cheap, without having
|
| 16 |
+
/// to write a SmallVector class) and contiguous only (so we can
|
| 17 |
+
/// return non-strided ArrayRef on index).
|
| 18 |
+
///
|
| 19 |
+
/// P.S. dimension 0 indexes rows, dimension 1 indexes columns
|
| 20 |
+
template <typename T>
|
| 21 |
+
class MatrixRef {
|
| 22 |
+
public:
|
| 23 |
+
typedef size_t size_type;
|
| 24 |
+
|
| 25 |
+
private:
|
| 26 |
+
/// Underlying ArrayRef
|
| 27 |
+
ArrayRef<T> arr;
|
| 28 |
+
|
| 29 |
+
/// Stride of dim 0 (outer dimension)
|
| 30 |
+
size_type stride0;
|
| 31 |
+
|
| 32 |
+
// Stride of dim 1 is assumed to be 1
|
| 33 |
+
|
| 34 |
+
public:
|
| 35 |
+
/// Construct an empty Matrixref.
|
| 36 |
+
/*implicit*/ MatrixRef() : arr(nullptr), stride0(0) {}
|
| 37 |
+
|
| 38 |
+
/// Construct an MatrixRef from an ArrayRef and outer stride.
|
| 39 |
+
/*implicit*/ MatrixRef(ArrayRef<T> arr, size_type stride0)
|
| 40 |
+
: arr(arr), stride0(stride0) {
|
| 41 |
+
TORCH_CHECK(
|
| 42 |
+
arr.size() % stride0 == 0,
|
| 43 |
+
"MatrixRef: ArrayRef size ",
|
| 44 |
+
arr.size(),
|
| 45 |
+
" not divisible by stride ",
|
| 46 |
+
stride0)
|
| 47 |
+
}
|
| 48 |
+
|
| 49 |
+
/// @}
|
| 50 |
+
/// @name Simple Operations
|
| 51 |
+
/// @{
|
| 52 |
+
|
| 53 |
+
/// empty - Check if the matrix is empty.
|
| 54 |
+
bool empty() const {
|
| 55 |
+
return arr.empty();
|
| 56 |
+
}
|
| 57 |
+
|
| 58 |
+
const T* data() const {
|
| 59 |
+
return arr.data();
|
| 60 |
+
}
|
| 61 |
+
|
| 62 |
+
/// size - Get size a dimension
|
| 63 |
+
size_t size(size_t dim) const {
|
| 64 |
+
if (dim == 0) {
|
| 65 |
+
return arr.size() / stride0;
|
| 66 |
+
} else if (dim == 1) {
|
| 67 |
+
return stride0;
|
| 68 |
+
} else {
|
| 69 |
+
TORCH_CHECK(
|
| 70 |
+
0, "MatrixRef: out of bounds dimension ", dim, "; expected 0 or 1");
|
| 71 |
+
}
|
| 72 |
+
}
|
| 73 |
+
|
| 74 |
+
size_t numel() const {
|
| 75 |
+
return arr.size();
|
| 76 |
+
}
|
| 77 |
+
|
| 78 |
+
/// equals - Check for element-wise equality.
|
| 79 |
+
bool equals(MatrixRef RHS) const {
|
| 80 |
+
return stride0 == RHS.stride0 && arr.equals(RHS.arr);
|
| 81 |
+
}
|
| 82 |
+
|
| 83 |
+
/// @}
|
| 84 |
+
/// @name Operator Overloads
|
| 85 |
+
/// @{
|
| 86 |
+
ArrayRef<T> operator[](size_t Index) const {
|
| 87 |
+
return arr.slice(Index * stride0, stride0);
|
| 88 |
+
}
|
| 89 |
+
|
| 90 |
+
/// Disallow accidental assignment from a temporary.
|
| 91 |
+
///
|
| 92 |
+
/// The declaration here is extra complicated so that "arrayRef = {}"
|
| 93 |
+
/// continues to select the move assignment operator.
|
| 94 |
+
template <typename U>
|
| 95 |
+
// NOLINTNEXTLINE(cppcoreguidelines-missing-std-forward)
|
| 96 |
+
std::enable_if_t<std::is_same_v<U, T>, MatrixRef<T>>& operator=(
|
| 97 |
+
U&& Temporary) = delete;
|
| 98 |
+
|
| 99 |
+
/// Disallow accidental assignment from a temporary.
|
| 100 |
+
///
|
| 101 |
+
/// The declaration here is extra complicated so that "arrayRef = {}"
|
| 102 |
+
/// continues to select the move assignment operator.
|
| 103 |
+
template <typename U>
|
| 104 |
+
std::enable_if_t<std::is_same_v<U, T>, MatrixRef<T>>& operator=(
|
| 105 |
+
std::initializer_list<U>) = delete;
|
| 106 |
+
};
|
| 107 |
+
|
| 108 |
+
} // end namespace at
|
lib/python3.12/site-packages/torch/include/ATen/MemoryOverlap.h
ADDED
|
@@ -0,0 +1,42 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <c10/macros/Export.h>
|
| 4 |
+
|
| 5 |
+
namespace c10 {
|
| 6 |
+
struct TensorImpl;
|
| 7 |
+
}
|
| 8 |
+
|
| 9 |
+
namespace at {
|
| 10 |
+
class TensorBase;
|
| 11 |
+
|
| 12 |
+
// MemOverlap: Whether or not there is memory overlap
|
| 13 |
+
//
|
| 14 |
+
// No: Absolutely no memory overlap
|
| 15 |
+
// Yes: Absolutely yes memory overlap
|
| 16 |
+
// TooHard: There might be memory overlap, but it was too expensive to compute.
|
| 17 |
+
//
|
| 18 |
+
// NB: Please update the python test for these if you renumber them.
|
| 19 |
+
enum class MemOverlap { No, Yes, TooHard };
|
| 20 |
+
|
| 21 |
+
enum class MemOverlapStatus { Full, Partial, No, TooHard };
|
| 22 |
+
|
| 23 |
+
TORCH_API MemOverlap has_internal_overlap(const TensorBase& t);
|
| 24 |
+
TORCH_API MemOverlap has_internal_overlap(c10::TensorImpl* t);
|
| 25 |
+
|
| 26 |
+
TORCH_API void assert_no_internal_overlap(const TensorBase& t);
|
| 27 |
+
TORCH_API void assert_no_internal_overlap(c10::TensorImpl* t);
|
| 28 |
+
|
| 29 |
+
TORCH_API MemOverlapStatus
|
| 30 |
+
get_overlap_status(const TensorBase& a, const TensorBase& b);
|
| 31 |
+
TORCH_API MemOverlapStatus
|
| 32 |
+
get_overlap_status(const c10::TensorImpl* a, const c10::TensorImpl* b);
|
| 33 |
+
|
| 34 |
+
TORCH_API void assert_no_partial_overlap(
|
| 35 |
+
const TensorBase& a,
|
| 36 |
+
const TensorBase& b);
|
| 37 |
+
void assert_no_partial_overlap(c10::TensorImpl* a, c10::TensorImpl* b);
|
| 38 |
+
|
| 39 |
+
TORCH_API void assert_no_overlap(const TensorBase& a, const TensorBase& b);
|
| 40 |
+
TORCH_API void assert_no_overlap(c10::TensorImpl* a, c10::TensorImpl* b);
|
| 41 |
+
|
| 42 |
+
} // namespace at
|
lib/python3.12/site-packages/torch/include/ATen/MetaFunctions.h
ADDED
|
@@ -0,0 +1,29 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#include <ATen/core/TensorBody.h>
|
| 2 |
+
|
| 3 |
+
// TODO Undo all logic introduced for Note [Avoiding Include Cycles In Static Dispatch]
|
| 4 |
+
// Code introduced to avoid cyclic dependency in static dispatch is no longer
|
| 5 |
+
// needed as static dispatch logic is moved from TensorBody.h, which caused cycles in the first place,
|
| 6 |
+
// to Operators.cpp for supporting multiple backends with multiple kernels.
|
| 7 |
+
//
|
| 8 |
+
// Note [Avoiding Include Cycles In Static Dispatch]
|
| 9 |
+
// In order to avoid #include cycles in the static dispatch build, we've carefully split out
|
| 10 |
+
// the static function definition files into {DispatchKey}Functions.h and {DispatchKey}Functions_inl.h.
|
| 11 |
+
//
|
| 12 |
+
// Without this split, the include cycle looks like TensorBody.h -> CPUFunctions.h -> TensorBody.h.
|
| 13 |
+
// - TensorBody.h #includes CPUFunctions.h in the static dispatch build, because the tensor methods
|
| 14 |
+
// all need to call into the fastpath C++ API defined in CPUFunctions.h. The methods are also all
|
| 15 |
+
// directly inlined into TensorBody.h.
|
| 16 |
+
// - CPUFunctions.h #includes TensorBody.h because it contains function declarations for the entire C++ API,
|
| 17 |
+
// which include functions that have defaultable std::optional<Tensor> arguments.
|
| 18 |
+
// That requires knowing the full Tensor class definition.
|
| 19 |
+
//
|
| 20 |
+
// We break the cycle by doing the following:
|
| 21 |
+
// - Split out CPUFunction.h into two files: CPUFunctions.h and CPUFunctions_inl.h
|
| 22 |
+
// - CPUFunction.h is a dummy file that just includes the Tensor class and includes CPUFunctions_inl.,
|
| 23 |
+
// - CPUFunctions_inl.h includes everything else
|
| 24 |
+
// - (only in the static dispatch build) TensorBody.h makes sure to finish defining the Tensor class,
|
| 25 |
+
// and then it includes CPUFunctions_inl.h.
|
| 26 |
+
// - All other files that want the cpu fastpath functions can include CPUFunctions.h directly.
|
| 27 |
+
// - This also means that static dispatch build, CPUFunctions.h only needs to
|
| 28 |
+
// #include TensorBody.h, and it will automatically bring in CPUFunctions_inl.h.
|
| 29 |
+
#include <ATen/MetaFunctions_inl.h>
|
lib/python3.12/site-packages/torch/include/ATen/NamedTensorUtils.h
ADDED
|
@@ -0,0 +1,214 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
#include <ATen/NamedTensor.h>
|
| 3 |
+
#include <ATen/TensorNames.h>
|
| 4 |
+
#include <ATen/WrapDimUtilsMulti.h>
|
| 5 |
+
|
| 6 |
+
#include <ATen/core/DimVector.h>
|
| 7 |
+
#include <ATen/core/Tensor.h>
|
| 8 |
+
|
| 9 |
+
namespace at {
|
| 10 |
+
|
| 11 |
+
using NameVector = SmallVector<Dimname, kDimVectorStaticSize>;
|
| 12 |
+
|
| 13 |
+
inline bool has_names(const ITensorListRef& tensors) {
|
| 14 |
+
return std::any_of(tensors.begin(), tensors.end(), [](const Tensor& t) {
|
| 15 |
+
return t.has_names();
|
| 16 |
+
});
|
| 17 |
+
}
|
| 18 |
+
|
| 19 |
+
// Converts dim to an positional index. Errors if `dim` cannot be used to
|
| 20 |
+
// refer to any dimension of tensor.
|
| 21 |
+
TORCH_API int64_t dimname_to_position(const Tensor& tensor, Dimname dim);
|
| 22 |
+
TORCH_API std::vector<int64_t> dimnames_to_positions(
|
| 23 |
+
const Tensor& tensor,
|
| 24 |
+
DimnameList dims);
|
| 25 |
+
|
| 26 |
+
// Unifies two DimnameList to produce a third. This is useful for implementing
|
| 27 |
+
// the named inference rule for binary broadcasting operations like add.
|
| 28 |
+
//
|
| 29 |
+
// There are three main constraints:
|
| 30 |
+
// 1) Check matching: Names must match positionally from the right.
|
| 31 |
+
// 2) Check misaligned: If a name `n` is in `names`, then it must appear at
|
| 32 |
+
// the same index from the right in other.
|
| 33 |
+
// 3) The output names are obtained by unifying the names individually from the
|
| 34 |
+
// right.
|
| 35 |
+
TORCH_API std::vector<Dimname> unify_from_right(
|
| 36 |
+
DimnameList names,
|
| 37 |
+
DimnameList other,
|
| 38 |
+
const char* action = "broadcast");
|
| 39 |
+
|
| 40 |
+
[[noreturn]] inline void reportNYIDimnameOverload(const char* op_name) {
|
| 41 |
+
TORCH_CHECK(
|
| 42 |
+
false,
|
| 43 |
+
op_name,
|
| 44 |
+
": You passed a dimname (string) to this op in place of a dimension "
|
| 45 |
+
"index but it does not yet support this behavior. Please pass a dimension "
|
| 46 |
+
"index to work around this.");
|
| 47 |
+
}
|
| 48 |
+
|
| 49 |
+
// [NOTE] Writing name inference rules
|
| 50 |
+
//
|
| 51 |
+
// Operators that support named tensors are either composed of operations that
|
| 52 |
+
// support named tensors or implement some name inference rule. An op that
|
| 53 |
+
// implements its own name inference rule generally looks like the following:
|
| 54 |
+
//
|
| 55 |
+
// Tensor op(...) {
|
| 56 |
+
// perform_shape_checks(...);
|
| 57 |
+
// # (1)
|
| 58 |
+
// auto maybe_outnames = compute_outnames(...);
|
| 59 |
+
// auto result = [&]() {
|
| 60 |
+
// NoNamesGuard guard;
|
| 61 |
+
// return op_impl(...);
|
| 62 |
+
// }();
|
| 63 |
+
// # (2)
|
| 64 |
+
// propagate_names_if_nonempty(result, maybe_outnames);
|
| 65 |
+
//
|
| 66 |
+
// Each op has (1) a compute outnames step and (2) a propagate names step.
|
| 67 |
+
//
|
| 68 |
+
// compute_outnames is responsible for checking that input names match and
|
| 69 |
+
// determining what the output names should be. It returns either:
|
| 70 |
+
// - {} (if the inputs tensors are all unnamed)
|
| 71 |
+
// - non-empty outnames.
|
| 72 |
+
//
|
| 73 |
+
// propagate_names_if_nonempty propagates the outnames if they exist to the
|
| 74 |
+
// result tensors.
|
| 75 |
+
//
|
| 76 |
+
// The {} case is an optimization; if the user does not use named tensors they
|
| 77 |
+
// pay no perf cost for it.
|
| 78 |
+
|
| 79 |
+
namespace namedinference {
|
| 80 |
+
|
| 81 |
+
const Tensor& propagate_names_if_present_and_nonempty(
|
| 82 |
+
const Tensor& result,
|
| 83 |
+
std::optional<DimnameList> maybe_names,
|
| 84 |
+
bool validate_names = false);
|
| 85 |
+
// Propagates `names` to `result` if `names` is not empty.
|
| 86 |
+
// `names` can be empty; see [NOTE] Writing name inference rules
|
| 87 |
+
// If `names` is not empty, `names.size()` should equal `result.dim()`.
|
| 88 |
+
// When in doubt, use this overload instead of the others.
|
| 89 |
+
TORCH_API const Tensor& propagate_names_if_nonempty(
|
| 90 |
+
const Tensor& result,
|
| 91 |
+
DimnameList maybe_names,
|
| 92 |
+
bool validate_names = false);
|
| 93 |
+
|
| 94 |
+
// Propagates `names` to `result`. Only use this if we are certain that there
|
| 95 |
+
// are names to propagate (that names is not empty).
|
| 96 |
+
TORCH_API const Tensor& propagate_names(
|
| 97 |
+
const Tensor& result,
|
| 98 |
+
DimnameList names,
|
| 99 |
+
bool validate_names = false);
|
| 100 |
+
|
| 101 |
+
// Propagates all names from src to result.
|
| 102 |
+
TORCH_API void propagate_names(const Tensor& result, const Tensor& src);
|
| 103 |
+
|
| 104 |
+
// Propagates all names except for those at the excluded_idxs.
|
| 105 |
+
TORCH_API void propagate_names_except(
|
| 106 |
+
const Tensor& result,
|
| 107 |
+
const Tensor& src,
|
| 108 |
+
IntArrayRef excluded_idxs);
|
| 109 |
+
|
| 110 |
+
// Used for reduction ops that have a `keepdim` arg.
|
| 111 |
+
TORCH_API void propagate_names_for_reduction(
|
| 112 |
+
const Tensor& result,
|
| 113 |
+
const Tensor& src,
|
| 114 |
+
IntArrayRef excluded_idxs,
|
| 115 |
+
bool keepdim);
|
| 116 |
+
|
| 117 |
+
TORCH_API void propagate_names_for_expand(
|
| 118 |
+
const Tensor& result,
|
| 119 |
+
const Tensor& self);
|
| 120 |
+
|
| 121 |
+
TORCH_API std::vector<Dimname> compute_cat_outnames(
|
| 122 |
+
const MaterializedITensorListRef& tensors);
|
| 123 |
+
|
| 124 |
+
TORCH_API std::vector<Dimname> compute_broadcast_outnames(
|
| 125 |
+
const Tensor& self,
|
| 126 |
+
const Tensor& other);
|
| 127 |
+
|
| 128 |
+
TORCH_API std::vector<Dimname> broadcast_to_outnames(
|
| 129 |
+
const Tensor& tensor,
|
| 130 |
+
const Tensor& reference_tensor,
|
| 131 |
+
const char* op_name);
|
| 132 |
+
|
| 133 |
+
TORCH_API std::vector<Dimname> compute_matmul_outnames(
|
| 134 |
+
const Tensor& self,
|
| 135 |
+
const Tensor& other);
|
| 136 |
+
|
| 137 |
+
TORCH_API std::vector<Dimname> compute_cdist_outnames(
|
| 138 |
+
const Tensor& self,
|
| 139 |
+
const Tensor& other);
|
| 140 |
+
|
| 141 |
+
TORCH_API std::vector<Dimname> compute_bmm_outnames(
|
| 142 |
+
const Tensor& result,
|
| 143 |
+
const Tensor& self,
|
| 144 |
+
const Tensor& other);
|
| 145 |
+
|
| 146 |
+
TORCH_API std::vector<Dimname> compute_squeeze_outnames(const Tensor& tensor);
|
| 147 |
+
TORCH_API std::vector<Dimname> compute_squeeze_outnames(
|
| 148 |
+
const Tensor& tensor,
|
| 149 |
+
std::bitset<dim_bitset_size> dims);
|
| 150 |
+
|
| 151 |
+
std::vector<Dimname> compute_diagonal_outnames(
|
| 152 |
+
const Tensor& tensor,
|
| 153 |
+
int64_t dim1,
|
| 154 |
+
int64_t dim2);
|
| 155 |
+
|
| 156 |
+
// TensorImpl* overloads for Legacy TH/THC code. Use these sparingly.
|
| 157 |
+
|
| 158 |
+
TORCH_API TensorImpl* propagate_names_if_nonempty(
|
| 159 |
+
TensorImpl* result,
|
| 160 |
+
DimnameList maybe_names,
|
| 161 |
+
bool validate_names = false);
|
| 162 |
+
|
| 163 |
+
TORCH_API TensorImpl* propagate_names(
|
| 164 |
+
TensorImpl* result,
|
| 165 |
+
DimnameList names,
|
| 166 |
+
bool validate_names = false);
|
| 167 |
+
|
| 168 |
+
TORCH_API void propagate_names(TensorImpl* result, /*const */ TensorImpl* src);
|
| 169 |
+
|
| 170 |
+
TORCH_API inline void propagate_names(
|
| 171 |
+
const TensorBase& result,
|
| 172 |
+
DimnameList names,
|
| 173 |
+
bool validate_names = false) {
|
| 174 |
+
propagate_names(result.unsafeGetTensorImpl(), names, validate_names);
|
| 175 |
+
}
|
| 176 |
+
|
| 177 |
+
TORCH_API inline void propagate_names_if_nonempty(
|
| 178 |
+
const TensorBase& result,
|
| 179 |
+
DimnameList names,
|
| 180 |
+
bool validate_names = false) {
|
| 181 |
+
propagate_names_if_nonempty(
|
| 182 |
+
result.unsafeGetTensorImpl(), names, validate_names);
|
| 183 |
+
}
|
| 184 |
+
|
| 185 |
+
TORCH_API inline void propagate_names(
|
| 186 |
+
const TensorBase& result,
|
| 187 |
+
const TensorBase& src) {
|
| 188 |
+
propagate_names(result.unsafeGetTensorImpl(), src.unsafeGetTensorImpl());
|
| 189 |
+
}
|
| 190 |
+
|
| 191 |
+
// result = m1 @ m2 + bias
|
| 192 |
+
TORCH_API std::vector<Dimname> propagate_names_for_addmm(
|
| 193 |
+
const Tensor& m1,
|
| 194 |
+
const Tensor& m2,
|
| 195 |
+
const Tensor& bias);
|
| 196 |
+
|
| 197 |
+
TORCH_API std::vector<Dimname> propagate_names_for_addmv(
|
| 198 |
+
const Tensor& mat,
|
| 199 |
+
const Tensor& vec,
|
| 200 |
+
const Tensor& bias);
|
| 201 |
+
|
| 202 |
+
TORCH_API void check_names_for_dot(TensorImpl* vec1, TensorImpl* vec2);
|
| 203 |
+
|
| 204 |
+
TORCH_API std::vector<Dimname> compute_baddbmm_outnames(
|
| 205 |
+
const Tensor& result,
|
| 206 |
+
const Tensor& self,
|
| 207 |
+
const Tensor& other,
|
| 208 |
+
const Tensor& bias);
|
| 209 |
+
|
| 210 |
+
TORCH_API bool are_names_equal(TensorImpl* self, TensorImpl* other);
|
| 211 |
+
|
| 212 |
+
} // namespace namedinference
|
| 213 |
+
|
| 214 |
+
} // namespace at
|
lib/python3.12/site-packages/torch/include/ATen/OpMathType.h
ADDED
|
@@ -0,0 +1,69 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <c10/core/ScalarType.h>
|
| 4 |
+
#include <c10/util/BFloat16.h>
|
| 5 |
+
#include <c10/util/Exception.h>
|
| 6 |
+
#include <c10/util/Float8_e4m3fn.h>
|
| 7 |
+
#include <c10/util/Float8_e4m3fnuz.h>
|
| 8 |
+
#include <c10/util/Float8_e5m2.h>
|
| 9 |
+
#include <c10/util/Float8_e5m2fnuz.h>
|
| 10 |
+
#include <c10/util/Half.h>
|
| 11 |
+
|
| 12 |
+
namespace at {
|
| 13 |
+
|
| 14 |
+
// For FP16 or BFloat16 inputs, ops should perform internal math in FP32.
|
| 15 |
+
template <typename scalar_t>
|
| 16 |
+
struct OpMathType {
|
| 17 |
+
using type = scalar_t;
|
| 18 |
+
};
|
| 19 |
+
template <>
|
| 20 |
+
struct OpMathType<at::Half> {
|
| 21 |
+
using type = float;
|
| 22 |
+
};
|
| 23 |
+
template <>
|
| 24 |
+
struct OpMathType<at::BFloat16> {
|
| 25 |
+
using type = float;
|
| 26 |
+
};
|
| 27 |
+
template <>
|
| 28 |
+
struct OpMathType<at::Float8_e5m2> {
|
| 29 |
+
using type = float;
|
| 30 |
+
};
|
| 31 |
+
template <>
|
| 32 |
+
struct OpMathType<at::Float8_e4m3fn> {
|
| 33 |
+
using type = float;
|
| 34 |
+
};
|
| 35 |
+
template <>
|
| 36 |
+
struct OpMathType<at::Float8_e5m2fnuz> {
|
| 37 |
+
using type = float;
|
| 38 |
+
};
|
| 39 |
+
template <>
|
| 40 |
+
struct OpMathType<at::Float8_e4m3fnuz> {
|
| 41 |
+
using type = float;
|
| 42 |
+
};
|
| 43 |
+
template <>
|
| 44 |
+
struct OpMathType<c10::complex<Half>> {
|
| 45 |
+
using type = c10::complex<float>;
|
| 46 |
+
};
|
| 47 |
+
|
| 48 |
+
template <typename T>
|
| 49 |
+
using opmath_type = typename OpMathType<T>::type;
|
| 50 |
+
|
| 51 |
+
namespace {
|
| 52 |
+
|
| 53 |
+
inline c10::ScalarType toOpMathType(const c10::ScalarType type) {
|
| 54 |
+
switch (type) {
|
| 55 |
+
#define DEFINE_CASE(scalar_t, TypeNum) \
|
| 56 |
+
case ScalarType::TypeNum: \
|
| 57 |
+
return CppTypeToScalarType<at::opmath_type<scalar_t>>::value;
|
| 58 |
+
|
| 59 |
+
AT_FORALL_SCALAR_TYPES_WITH_COMPLEX(DEFINE_CASE)
|
| 60 |
+
#undef DEFINE_CASE
|
| 61 |
+
|
| 62 |
+
default:
|
| 63 |
+
TORCH_INTERNAL_ASSERT(false, "Unrecognized ScalarType: ", type);
|
| 64 |
+
}
|
| 65 |
+
}
|
| 66 |
+
|
| 67 |
+
} // namespace
|
| 68 |
+
|
| 69 |
+
} // namespace at
|
lib/python3.12/site-packages/torch/include/ATen/Operators.h
ADDED
|
@@ -0,0 +1,1389 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
// @generated by torchgen/gen.py from Operators.h
|
| 4 |
+
|
| 5 |
+
#ifdef TORCH_ASSERT_NO_OPERATORS
|
| 6 |
+
#error This change adds a dependency on native_functions.yaml, \
|
| 7 |
+
meaning the file will need to be re-compiled every time an operator \
|
| 8 |
+
is changed or added. Consider if your change would be better placed in \
|
| 9 |
+
another file, or if a more specific header might achieve the same goal. \
|
| 10 |
+
See NOTE: [Tensor vs. TensorBase]
|
| 11 |
+
#endif
|
| 12 |
+
|
| 13 |
+
#if defined(AT_PER_OPERATOR_HEADERS) && defined(TORCH_ASSERT_ONLY_METHOD_OPERATORS)
|
| 14 |
+
#error This change adds a dependency on all pytorch operators, meaning the \
|
| 15 |
+
file will need to be re-compiled every time an operator is changed or added. \
|
| 16 |
+
Consider including a specific operator from <ATen/ops/{my_operator}_ops.h> \
|
| 17 |
+
and see NOTE [TORCH_ASSERT_ONLY_METHOD_OPERATORS].
|
| 18 |
+
#endif
|
| 19 |
+
|
| 20 |
+
#include <c10/core/SymInt.h>
|
| 21 |
+
#include <c10/core/SymIntArrayRef.h>
|
| 22 |
+
#include <c10/core/Scalar.h>
|
| 23 |
+
#include <c10/core/TensorOptions.h>
|
| 24 |
+
#include <c10/core/QScheme.h>
|
| 25 |
+
#include <c10/util/OptionalArrayRef.h>
|
| 26 |
+
#include <tuple>
|
| 27 |
+
#include <vector>
|
| 28 |
+
|
| 29 |
+
#include <ATen/ops/_adaptive_avg_pool2d_ops.h>
|
| 30 |
+
#include <ATen/ops/_adaptive_avg_pool2d_backward_ops.h>
|
| 31 |
+
#include <ATen/ops/_adaptive_avg_pool3d_ops.h>
|
| 32 |
+
#include <ATen/ops/_adaptive_avg_pool3d_backward_ops.h>
|
| 33 |
+
#include <ATen/ops/_add_batch_dim_ops.h>
|
| 34 |
+
#include <ATen/ops/_add_relu_ops.h>
|
| 35 |
+
#include <ATen/ops/_addmm_activation_ops.h>
|
| 36 |
+
#include <ATen/ops/_aminmax_ops.h>
|
| 37 |
+
#include <ATen/ops/_amp_foreach_non_finite_check_and_unscale_ops.h>
|
| 38 |
+
#include <ATen/ops/_amp_update_scale_ops.h>
|
| 39 |
+
#include <ATen/ops/_assert_async_ops.h>
|
| 40 |
+
#include <ATen/ops/_assert_scalar_ops.h>
|
| 41 |
+
#include <ATen/ops/_assert_tensor_metadata_ops.h>
|
| 42 |
+
#include <ATen/ops/_autocast_to_full_precision_ops.h>
|
| 43 |
+
#include <ATen/ops/_autocast_to_reduced_precision_ops.h>
|
| 44 |
+
#include <ATen/ops/_backward_ops.h>
|
| 45 |
+
#include <ATen/ops/_batch_norm_impl_index_ops.h>
|
| 46 |
+
#include <ATen/ops/_batch_norm_impl_index_backward_ops.h>
|
| 47 |
+
#include <ATen/ops/_batch_norm_no_update_ops.h>
|
| 48 |
+
#include <ATen/ops/_batch_norm_with_update_ops.h>
|
| 49 |
+
#include <ATen/ops/_cast_Byte_ops.h>
|
| 50 |
+
#include <ATen/ops/_cast_Char_ops.h>
|
| 51 |
+
#include <ATen/ops/_cast_Double_ops.h>
|
| 52 |
+
#include <ATen/ops/_cast_Float_ops.h>
|
| 53 |
+
#include <ATen/ops/_cast_Half_ops.h>
|
| 54 |
+
#include <ATen/ops/_cast_Int_ops.h>
|
| 55 |
+
#include <ATen/ops/_cast_Long_ops.h>
|
| 56 |
+
#include <ATen/ops/_cast_Short_ops.h>
|
| 57 |
+
#include <ATen/ops/_cdist_backward_ops.h>
|
| 58 |
+
#include <ATen/ops/_cdist_forward_ops.h>
|
| 59 |
+
#include <ATen/ops/_cholesky_solve_helper_ops.h>
|
| 60 |
+
#include <ATen/ops/_choose_qparams_per_tensor_ops.h>
|
| 61 |
+
#include <ATen/ops/_chunk_cat_ops.h>
|
| 62 |
+
#include <ATen/ops/_coalesce_ops.h>
|
| 63 |
+
#include <ATen/ops/_coalesced_ops.h>
|
| 64 |
+
#include <ATen/ops/_compute_linear_combination_ops.h>
|
| 65 |
+
#include <ATen/ops/_conj_ops.h>
|
| 66 |
+
#include <ATen/ops/_conj_copy_ops.h>
|
| 67 |
+
#include <ATen/ops/_conj_physical_ops.h>
|
| 68 |
+
#include <ATen/ops/_conv_depthwise2d_ops.h>
|
| 69 |
+
#include <ATen/ops/_convert_indices_from_coo_to_csr_ops.h>
|
| 70 |
+
#include <ATen/ops/_convert_indices_from_csr_to_coo_ops.h>
|
| 71 |
+
#include <ATen/ops/_convert_weight_to_int4pack_ops.h>
|
| 72 |
+
#include <ATen/ops/_convert_weight_to_int4pack_for_cpu_ops.h>
|
| 73 |
+
#include <ATen/ops/_convolution_ops.h>
|
| 74 |
+
#include <ATen/ops/_convolution_double_backward_ops.h>
|
| 75 |
+
#include <ATen/ops/_convolution_mode_ops.h>
|
| 76 |
+
#include <ATen/ops/_copy_from_ops.h>
|
| 77 |
+
#include <ATen/ops/_copy_from_and_resize_ops.h>
|
| 78 |
+
#include <ATen/ops/_cslt_compress_ops.h>
|
| 79 |
+
#include <ATen/ops/_cslt_sparse_mm_ops.h>
|
| 80 |
+
#include <ATen/ops/_cslt_sparse_mm_search_ops.h>
|
| 81 |
+
#include <ATen/ops/_ctc_loss_ops.h>
|
| 82 |
+
#include <ATen/ops/_ctc_loss_backward_ops.h>
|
| 83 |
+
#include <ATen/ops/_cudnn_ctc_loss_ops.h>
|
| 84 |
+
#include <ATen/ops/_cudnn_init_dropout_state_ops.h>
|
| 85 |
+
#include <ATen/ops/_cudnn_rnn_ops.h>
|
| 86 |
+
#include <ATen/ops/_cudnn_rnn_backward_ops.h>
|
| 87 |
+
#include <ATen/ops/_cudnn_rnn_flatten_weight_ops.h>
|
| 88 |
+
#include <ATen/ops/_cufft_clear_plan_cache_ops.h>
|
| 89 |
+
#include <ATen/ops/_cufft_get_plan_cache_max_size_ops.h>
|
| 90 |
+
#include <ATen/ops/_cufft_get_plan_cache_size_ops.h>
|
| 91 |
+
#include <ATen/ops/_cufft_set_plan_cache_max_size_ops.h>
|
| 92 |
+
#include <ATen/ops/_cummax_helper_ops.h>
|
| 93 |
+
#include <ATen/ops/_cummin_helper_ops.h>
|
| 94 |
+
#include <ATen/ops/_debug_has_internal_overlap_ops.h>
|
| 95 |
+
#include <ATen/ops/_dimI_ops.h>
|
| 96 |
+
#include <ATen/ops/_dimV_ops.h>
|
| 97 |
+
#include <ATen/ops/_dim_arange_ops.h>
|
| 98 |
+
#include <ATen/ops/_dirichlet_grad_ops.h>
|
| 99 |
+
#include <ATen/ops/_efficient_attention_backward_ops.h>
|
| 100 |
+
#include <ATen/ops/_efficient_attention_forward_ops.h>
|
| 101 |
+
#include <ATen/ops/_efficientzerotensor_ops.h>
|
| 102 |
+
#include <ATen/ops/_embedding_bag_ops.h>
|
| 103 |
+
#include <ATen/ops/_embedding_bag_backward_ops.h>
|
| 104 |
+
#include <ATen/ops/_embedding_bag_dense_backward_ops.h>
|
| 105 |
+
#include <ATen/ops/_embedding_bag_forward_only_ops.h>
|
| 106 |
+
#include <ATen/ops/_embedding_bag_per_sample_weights_backward_ops.h>
|
| 107 |
+
#include <ATen/ops/_embedding_bag_sparse_backward_ops.h>
|
| 108 |
+
#include <ATen/ops/_empty_affine_quantized_ops.h>
|
| 109 |
+
#include <ATen/ops/_empty_per_channel_affine_quantized_ops.h>
|
| 110 |
+
#include <ATen/ops/_euclidean_dist_ops.h>
|
| 111 |
+
#include <ATen/ops/_fake_quantize_learnable_per_channel_affine_ops.h>
|
| 112 |
+
#include <ATen/ops/_fake_quantize_learnable_per_channel_affine_backward_ops.h>
|
| 113 |
+
#include <ATen/ops/_fake_quantize_learnable_per_tensor_affine_ops.h>
|
| 114 |
+
#include <ATen/ops/_fake_quantize_learnable_per_tensor_affine_backward_ops.h>
|
| 115 |
+
#include <ATen/ops/_fake_quantize_per_tensor_affine_cachemask_tensor_qparams_ops.h>
|
| 116 |
+
#include <ATen/ops/_fft_c2c_ops.h>
|
| 117 |
+
#include <ATen/ops/_fft_c2r_ops.h>
|
| 118 |
+
#include <ATen/ops/_fft_r2c_ops.h>
|
| 119 |
+
#include <ATen/ops/_fill_mem_eff_dropout_mask_ops.h>
|
| 120 |
+
#include <ATen/ops/_flash_attention_backward_ops.h>
|
| 121 |
+
#include <ATen/ops/_flash_attention_forward_ops.h>
|
| 122 |
+
#include <ATen/ops/_foobar_ops.h>
|
| 123 |
+
#include <ATen/ops/_foreach_abs_ops.h>
|
| 124 |
+
#include <ATen/ops/_foreach_acos_ops.h>
|
| 125 |
+
#include <ATen/ops/_foreach_add_ops.h>
|
| 126 |
+
#include <ATen/ops/_foreach_addcdiv_ops.h>
|
| 127 |
+
#include <ATen/ops/_foreach_addcmul_ops.h>
|
| 128 |
+
#include <ATen/ops/_foreach_asin_ops.h>
|
| 129 |
+
#include <ATen/ops/_foreach_atan_ops.h>
|
| 130 |
+
#include <ATen/ops/_foreach_ceil_ops.h>
|
| 131 |
+
#include <ATen/ops/_foreach_clamp_max_ops.h>
|
| 132 |
+
#include <ATen/ops/_foreach_clamp_min_ops.h>
|
| 133 |
+
#include <ATen/ops/_foreach_copy_ops.h>
|
| 134 |
+
#include <ATen/ops/_foreach_cos_ops.h>
|
| 135 |
+
#include <ATen/ops/_foreach_cosh_ops.h>
|
| 136 |
+
#include <ATen/ops/_foreach_div_ops.h>
|
| 137 |
+
#include <ATen/ops/_foreach_erf_ops.h>
|
| 138 |
+
#include <ATen/ops/_foreach_erfc_ops.h>
|
| 139 |
+
#include <ATen/ops/_foreach_exp_ops.h>
|
| 140 |
+
#include <ATen/ops/_foreach_expm1_ops.h>
|
| 141 |
+
#include <ATen/ops/_foreach_floor_ops.h>
|
| 142 |
+
#include <ATen/ops/_foreach_frac_ops.h>
|
| 143 |
+
#include <ATen/ops/_foreach_lerp_ops.h>
|
| 144 |
+
#include <ATen/ops/_foreach_lgamma_ops.h>
|
| 145 |
+
#include <ATen/ops/_foreach_log_ops.h>
|
| 146 |
+
#include <ATen/ops/_foreach_log10_ops.h>
|
| 147 |
+
#include <ATen/ops/_foreach_log1p_ops.h>
|
| 148 |
+
#include <ATen/ops/_foreach_log2_ops.h>
|
| 149 |
+
#include <ATen/ops/_foreach_max_ops.h>
|
| 150 |
+
#include <ATen/ops/_foreach_maximum_ops.h>
|
| 151 |
+
#include <ATen/ops/_foreach_minimum_ops.h>
|
| 152 |
+
#include <ATen/ops/_foreach_mul_ops.h>
|
| 153 |
+
#include <ATen/ops/_foreach_neg_ops.h>
|
| 154 |
+
#include <ATen/ops/_foreach_norm_ops.h>
|
| 155 |
+
#include <ATen/ops/_foreach_pow_ops.h>
|
| 156 |
+
#include <ATen/ops/_foreach_reciprocal_ops.h>
|
| 157 |
+
#include <ATen/ops/_foreach_round_ops.h>
|
| 158 |
+
#include <ATen/ops/_foreach_rsqrt_ops.h>
|
| 159 |
+
#include <ATen/ops/_foreach_sigmoid_ops.h>
|
| 160 |
+
#include <ATen/ops/_foreach_sign_ops.h>
|
| 161 |
+
#include <ATen/ops/_foreach_sin_ops.h>
|
| 162 |
+
#include <ATen/ops/_foreach_sinh_ops.h>
|
| 163 |
+
#include <ATen/ops/_foreach_sqrt_ops.h>
|
| 164 |
+
#include <ATen/ops/_foreach_sub_ops.h>
|
| 165 |
+
#include <ATen/ops/_foreach_tan_ops.h>
|
| 166 |
+
#include <ATen/ops/_foreach_tanh_ops.h>
|
| 167 |
+
#include <ATen/ops/_foreach_trunc_ops.h>
|
| 168 |
+
#include <ATen/ops/_foreach_zero_ops.h>
|
| 169 |
+
#include <ATen/ops/_functional_assert_async_ops.h>
|
| 170 |
+
#include <ATen/ops/_functional_assert_scalar_ops.h>
|
| 171 |
+
#include <ATen/ops/_functional_sym_constrain_range_ops.h>
|
| 172 |
+
#include <ATen/ops/_functional_sym_constrain_range_for_size_ops.h>
|
| 173 |
+
#include <ATen/ops/_fused_adagrad_ops.h>
|
| 174 |
+
#include <ATen/ops/_fused_adam_ops.h>
|
| 175 |
+
#include <ATen/ops/_fused_adamw_ops.h>
|
| 176 |
+
#include <ATen/ops/_fused_dropout_ops.h>
|
| 177 |
+
#include <ATen/ops/_fused_moving_avg_obs_fq_helper_ops.h>
|
| 178 |
+
#include <ATen/ops/_fused_sdp_choice_ops.h>
|
| 179 |
+
#include <ATen/ops/_fused_sgd_ops.h>
|
| 180 |
+
#include <ATen/ops/_fw_primal_ops.h>
|
| 181 |
+
#include <ATen/ops/_fw_primal_copy_ops.h>
|
| 182 |
+
#include <ATen/ops/_gather_sparse_backward_ops.h>
|
| 183 |
+
#include <ATen/ops/_grid_sampler_2d_cpu_fallback_ops.h>
|
| 184 |
+
#include <ATen/ops/_grid_sampler_2d_cpu_fallback_backward_ops.h>
|
| 185 |
+
#include <ATen/ops/_has_compatible_shallow_copy_type_ops.h>
|
| 186 |
+
#include <ATen/ops/_has_same_storage_numel_ops.h>
|
| 187 |
+
#include <ATen/ops/_histogramdd_bin_edges_ops.h>
|
| 188 |
+
#include <ATen/ops/_histogramdd_from_bin_cts_ops.h>
|
| 189 |
+
#include <ATen/ops/_histogramdd_from_bin_tensors_ops.h>
|
| 190 |
+
#include <ATen/ops/_index_put_impl_ops.h>
|
| 191 |
+
#include <ATen/ops/_indices_ops.h>
|
| 192 |
+
#include <ATen/ops/_indices_copy_ops.h>
|
| 193 |
+
#include <ATen/ops/_int_mm_ops.h>
|
| 194 |
+
#include <ATen/ops/_is_all_true_ops.h>
|
| 195 |
+
#include <ATen/ops/_is_any_true_ops.h>
|
| 196 |
+
#include <ATen/ops/_is_zerotensor_ops.h>
|
| 197 |
+
#include <ATen/ops/_jagged_to_padded_dense_forward_ops.h>
|
| 198 |
+
#include <ATen/ops/_lazy_clone_ops.h>
|
| 199 |
+
#include <ATen/ops/_linalg_check_errors_ops.h>
|
| 200 |
+
#include <ATen/ops/_linalg_det_ops.h>
|
| 201 |
+
#include <ATen/ops/_linalg_eigh_ops.h>
|
| 202 |
+
#include <ATen/ops/_linalg_eigvals_ops.h>
|
| 203 |
+
#include <ATen/ops/_linalg_slogdet_ops.h>
|
| 204 |
+
#include <ATen/ops/_linalg_solve_ex_ops.h>
|
| 205 |
+
#include <ATen/ops/_linalg_svd_ops.h>
|
| 206 |
+
#include <ATen/ops/_local_scalar_dense_ops.h>
|
| 207 |
+
#include <ATen/ops/_log_softmax_ops.h>
|
| 208 |
+
#include <ATen/ops/_log_softmax_backward_data_ops.h>
|
| 209 |
+
#include <ATen/ops/_logcumsumexp_ops.h>
|
| 210 |
+
#include <ATen/ops/_lstm_mps_ops.h>
|
| 211 |
+
#include <ATen/ops/_lu_with_info_ops.h>
|
| 212 |
+
#include <ATen/ops/_make_dep_token_ops.h>
|
| 213 |
+
#include <ATen/ops/_make_dual_ops.h>
|
| 214 |
+
#include <ATen/ops/_make_dual_copy_ops.h>
|
| 215 |
+
#include <ATen/ops/_make_per_channel_quantized_tensor_ops.h>
|
| 216 |
+
#include <ATen/ops/_make_per_tensor_quantized_tensor_ops.h>
|
| 217 |
+
#include <ATen/ops/_masked_scale_ops.h>
|
| 218 |
+
#include <ATen/ops/_masked_softmax_ops.h>
|
| 219 |
+
#include <ATen/ops/_masked_softmax_backward_ops.h>
|
| 220 |
+
#include <ATen/ops/_mixed_dtypes_linear_ops.h>
|
| 221 |
+
#include <ATen/ops/_mkldnn_reshape_ops.h>
|
| 222 |
+
#include <ATen/ops/_mkldnn_transpose_ops.h>
|
| 223 |
+
#include <ATen/ops/_mps_convolution_ops.h>
|
| 224 |
+
#include <ATen/ops/_mps_convolution_transpose_ops.h>
|
| 225 |
+
#include <ATen/ops/_native_batch_norm_legit_ops.h>
|
| 226 |
+
#include <ATen/ops/_native_batch_norm_legit_no_training_ops.h>
|
| 227 |
+
#include <ATen/ops/_native_multi_head_attention_ops.h>
|
| 228 |
+
#include <ATen/ops/_neg_view_ops.h>
|
| 229 |
+
#include <ATen/ops/_neg_view_copy_ops.h>
|
| 230 |
+
#include <ATen/ops/_nested_compute_contiguous_strides_offsets_ops.h>
|
| 231 |
+
#include <ATen/ops/_nested_from_padded_ops.h>
|
| 232 |
+
#include <ATen/ops/_nested_from_padded_and_nested_example_ops.h>
|
| 233 |
+
#include <ATen/ops/_nested_from_padded_tensor_ops.h>
|
| 234 |
+
#include <ATen/ops/_nested_get_jagged_dummy_ops.h>
|
| 235 |
+
#include <ATen/ops/_nested_get_lengths_ops.h>
|
| 236 |
+
#include <ATen/ops/_nested_get_max_seqlen_ops.h>
|
| 237 |
+
#include <ATen/ops/_nested_get_min_seqlen_ops.h>
|
| 238 |
+
#include <ATen/ops/_nested_get_offsets_ops.h>
|
| 239 |
+
#include <ATen/ops/_nested_get_ragged_idx_ops.h>
|
| 240 |
+
#include <ATen/ops/_nested_get_values_ops.h>
|
| 241 |
+
#include <ATen/ops/_nested_get_values_copy_ops.h>
|
| 242 |
+
#include <ATen/ops/_nested_select_backward_ops.h>
|
| 243 |
+
#include <ATen/ops/_nested_sum_backward_ops.h>
|
| 244 |
+
#include <ATen/ops/_nested_tensor_from_mask_ops.h>
|
| 245 |
+
#include <ATen/ops/_nested_tensor_from_mask_left_aligned_ops.h>
|
| 246 |
+
#include <ATen/ops/_nested_tensor_from_tensor_list_ops.h>
|
| 247 |
+
#include <ATen/ops/_nested_tensor_size_ops.h>
|
| 248 |
+
#include <ATen/ops/_nested_tensor_softmax_with_shape_ops.h>
|
| 249 |
+
#include <ATen/ops/_nested_tensor_storage_offsets_ops.h>
|
| 250 |
+
#include <ATen/ops/_nested_tensor_strides_ops.h>
|
| 251 |
+
#include <ATen/ops/_nested_view_from_buffer_ops.h>
|
| 252 |
+
#include <ATen/ops/_nested_view_from_buffer_copy_ops.h>
|
| 253 |
+
#include <ATen/ops/_nested_view_from_jagged_ops.h>
|
| 254 |
+
#include <ATen/ops/_nested_view_from_jagged_copy_ops.h>
|
| 255 |
+
#include <ATen/ops/_new_zeros_with_same_feature_meta_ops.h>
|
| 256 |
+
#include <ATen/ops/_nnpack_available_ops.h>
|
| 257 |
+
#include <ATen/ops/_nnpack_spatial_convolution_ops.h>
|
| 258 |
+
#include <ATen/ops/_nnz_ops.h>
|
| 259 |
+
#include <ATen/ops/_pack_padded_sequence_ops.h>
|
| 260 |
+
#include <ATen/ops/_pack_padded_sequence_backward_ops.h>
|
| 261 |
+
#include <ATen/ops/_pad_circular_ops.h>
|
| 262 |
+
#include <ATen/ops/_pad_enum_ops.h>
|
| 263 |
+
#include <ATen/ops/_pad_packed_sequence_ops.h>
|
| 264 |
+
#include <ATen/ops/_padded_dense_to_jagged_forward_ops.h>
|
| 265 |
+
#include <ATen/ops/_pdist_backward_ops.h>
|
| 266 |
+
#include <ATen/ops/_pdist_forward_ops.h>
|
| 267 |
+
#include <ATen/ops/_pin_memory_ops.h>
|
| 268 |
+
#include <ATen/ops/_prelu_kernel_ops.h>
|
| 269 |
+
#include <ATen/ops/_prelu_kernel_backward_ops.h>
|
| 270 |
+
#include <ATen/ops/_print_ops.h>
|
| 271 |
+
#include <ATen/ops/_propagate_xla_data_ops.h>
|
| 272 |
+
#include <ATen/ops/_remove_batch_dim_ops.h>
|
| 273 |
+
#include <ATen/ops/_reshape_alias_ops.h>
|
| 274 |
+
#include <ATen/ops/_reshape_alias_copy_ops.h>
|
| 275 |
+
#include <ATen/ops/_reshape_copy_ops.h>
|
| 276 |
+
#include <ATen/ops/_reshape_from_tensor_ops.h>
|
| 277 |
+
#include <ATen/ops/_resize_output_ops.h>
|
| 278 |
+
#include <ATen/ops/_rowwise_prune_ops.h>
|
| 279 |
+
#include <ATen/ops/_safe_softmax_ops.h>
|
| 280 |
+
#include <ATen/ops/_sample_dirichlet_ops.h>
|
| 281 |
+
#include <ATen/ops/_saturate_weight_to_fp16_ops.h>
|
| 282 |
+
#include <ATen/ops/_scaled_dot_product_attention_math_ops.h>
|
| 283 |
+
#include <ATen/ops/_scaled_dot_product_attention_math_for_mps_ops.h>
|
| 284 |
+
#include <ATen/ops/_scaled_dot_product_cudnn_attention_ops.h>
|
| 285 |
+
#include <ATen/ops/_scaled_dot_product_cudnn_attention_backward_ops.h>
|
| 286 |
+
#include <ATen/ops/_scaled_dot_product_efficient_attention_ops.h>
|
| 287 |
+
#include <ATen/ops/_scaled_dot_product_efficient_attention_backward_ops.h>
|
| 288 |
+
#include <ATen/ops/_scaled_dot_product_flash_attention_ops.h>
|
| 289 |
+
#include <ATen/ops/_scaled_dot_product_flash_attention_backward_ops.h>
|
| 290 |
+
#include <ATen/ops/_scaled_dot_product_flash_attention_for_cpu_ops.h>
|
| 291 |
+
#include <ATen/ops/_scaled_dot_product_flash_attention_for_cpu_backward_ops.h>
|
| 292 |
+
#include <ATen/ops/_scaled_dot_product_fused_attention_overrideable_ops.h>
|
| 293 |
+
#include <ATen/ops/_scaled_dot_product_fused_attention_overrideable_backward_ops.h>
|
| 294 |
+
#include <ATen/ops/_scaled_mm_ops.h>
|
| 295 |
+
#include <ATen/ops/_segment_reduce_backward_ops.h>
|
| 296 |
+
#include <ATen/ops/_shape_as_tensor_ops.h>
|
| 297 |
+
#include <ATen/ops/_slow_conv2d_backward_ops.h>
|
| 298 |
+
#include <ATen/ops/_slow_conv2d_forward_ops.h>
|
| 299 |
+
#include <ATen/ops/_sobol_engine_draw_ops.h>
|
| 300 |
+
#include <ATen/ops/_sobol_engine_ff_ops.h>
|
| 301 |
+
#include <ATen/ops/_sobol_engine_initialize_state_ops.h>
|
| 302 |
+
#include <ATen/ops/_sobol_engine_scramble_ops.h>
|
| 303 |
+
#include <ATen/ops/_softmax_ops.h>
|
| 304 |
+
#include <ATen/ops/_softmax_backward_data_ops.h>
|
| 305 |
+
#include <ATen/ops/_sparse_addmm_ops.h>
|
| 306 |
+
#include <ATen/ops/_sparse_broadcast_to_ops.h>
|
| 307 |
+
#include <ATen/ops/_sparse_broadcast_to_copy_ops.h>
|
| 308 |
+
#include <ATen/ops/_sparse_bsc_tensor_unsafe_ops.h>
|
| 309 |
+
#include <ATen/ops/_sparse_bsr_tensor_unsafe_ops.h>
|
| 310 |
+
#include <ATen/ops/_sparse_compressed_tensor_unsafe_ops.h>
|
| 311 |
+
#include <ATen/ops/_sparse_compressed_tensor_with_dims_ops.h>
|
| 312 |
+
#include <ATen/ops/_sparse_coo_tensor_unsafe_ops.h>
|
| 313 |
+
#include <ATen/ops/_sparse_coo_tensor_with_dims_ops.h>
|
| 314 |
+
#include <ATen/ops/_sparse_coo_tensor_with_dims_and_tensors_ops.h>
|
| 315 |
+
#include <ATen/ops/_sparse_csc_tensor_unsafe_ops.h>
|
| 316 |
+
#include <ATen/ops/_sparse_csr_prod_ops.h>
|
| 317 |
+
#include <ATen/ops/_sparse_csr_sum_ops.h>
|
| 318 |
+
#include <ATen/ops/_sparse_csr_tensor_unsafe_ops.h>
|
| 319 |
+
#include <ATen/ops/_sparse_log_softmax_ops.h>
|
| 320 |
+
#include <ATen/ops/_sparse_log_softmax_backward_data_ops.h>
|
| 321 |
+
#include <ATen/ops/_sparse_mask_projection_ops.h>
|
| 322 |
+
#include <ATen/ops/_sparse_mm_ops.h>
|
| 323 |
+
#include <ATen/ops/_sparse_mm_reduce_impl_ops.h>
|
| 324 |
+
#include <ATen/ops/_sparse_mm_reduce_impl_backward_ops.h>
|
| 325 |
+
#include <ATen/ops/_sparse_semi_structured_addmm_ops.h>
|
| 326 |
+
#include <ATen/ops/_sparse_semi_structured_apply_ops.h>
|
| 327 |
+
#include <ATen/ops/_sparse_semi_structured_apply_dense_ops.h>
|
| 328 |
+
#include <ATen/ops/_sparse_semi_structured_linear_ops.h>
|
| 329 |
+
#include <ATen/ops/_sparse_semi_structured_mm_ops.h>
|
| 330 |
+
#include <ATen/ops/_sparse_semi_structured_tile_ops.h>
|
| 331 |
+
#include <ATen/ops/_sparse_softmax_ops.h>
|
| 332 |
+
#include <ATen/ops/_sparse_softmax_backward_data_ops.h>
|
| 333 |
+
#include <ATen/ops/_sparse_sparse_matmul_ops.h>
|
| 334 |
+
#include <ATen/ops/_sparse_sum_ops.h>
|
| 335 |
+
#include <ATen/ops/_sparse_sum_backward_ops.h>
|
| 336 |
+
#include <ATen/ops/_spdiags_ops.h>
|
| 337 |
+
#include <ATen/ops/_spsolve_ops.h>
|
| 338 |
+
#include <ATen/ops/_stack_ops.h>
|
| 339 |
+
#include <ATen/ops/_standard_gamma_ops.h>
|
| 340 |
+
#include <ATen/ops/_standard_gamma_grad_ops.h>
|
| 341 |
+
#include <ATen/ops/_test_ambiguous_defaults_ops.h>
|
| 342 |
+
#include <ATen/ops/_test_autograd_multiple_dispatch_ops.h>
|
| 343 |
+
#include <ATen/ops/_test_autograd_multiple_dispatch_view_ops.h>
|
| 344 |
+
#include <ATen/ops/_test_autograd_multiple_dispatch_view_copy_ops.h>
|
| 345 |
+
#include <ATen/ops/_test_check_tensor_ops.h>
|
| 346 |
+
#include <ATen/ops/_test_functorch_fallback_ops.h>
|
| 347 |
+
#include <ATen/ops/_test_optional_filled_intlist_ops.h>
|
| 348 |
+
#include <ATen/ops/_test_optional_floatlist_ops.h>
|
| 349 |
+
#include <ATen/ops/_test_optional_intlist_ops.h>
|
| 350 |
+
#include <ATen/ops/_test_parallel_materialize_ops.h>
|
| 351 |
+
#include <ATen/ops/_test_serialization_subcmul_ops.h>
|
| 352 |
+
#include <ATen/ops/_test_string_default_ops.h>
|
| 353 |
+
#include <ATen/ops/_test_warn_in_autograd_ops.h>
|
| 354 |
+
#include <ATen/ops/_thnn_differentiable_gru_cell_backward_ops.h>
|
| 355 |
+
#include <ATen/ops/_thnn_differentiable_lstm_cell_backward_ops.h>
|
| 356 |
+
#include <ATen/ops/_thnn_fused_gru_cell_ops.h>
|
| 357 |
+
#include <ATen/ops/_thnn_fused_gru_cell_backward_ops.h>
|
| 358 |
+
#include <ATen/ops/_thnn_fused_lstm_cell_ops.h>
|
| 359 |
+
#include <ATen/ops/_thnn_fused_lstm_cell_backward_ops.h>
|
| 360 |
+
#include <ATen/ops/_thnn_fused_lstm_cell_backward_impl_ops.h>
|
| 361 |
+
#include <ATen/ops/_to_copy_ops.h>
|
| 362 |
+
#include <ATen/ops/_to_cpu_ops.h>
|
| 363 |
+
#include <ATen/ops/_to_dense_ops.h>
|
| 364 |
+
#include <ATen/ops/_to_sparse_ops.h>
|
| 365 |
+
#include <ATen/ops/_to_sparse_bsc_ops.h>
|
| 366 |
+
#include <ATen/ops/_to_sparse_bsr_ops.h>
|
| 367 |
+
#include <ATen/ops/_to_sparse_csc_ops.h>
|
| 368 |
+
#include <ATen/ops/_to_sparse_csr_ops.h>
|
| 369 |
+
#include <ATen/ops/_to_sparse_semi_structured_ops.h>
|
| 370 |
+
#include <ATen/ops/_transform_bias_rescale_qkv_ops.h>
|
| 371 |
+
#include <ATen/ops/_transformer_encoder_layer_fwd_ops.h>
|
| 372 |
+
#include <ATen/ops/_trilinear_ops.h>
|
| 373 |
+
#include <ATen/ops/_triton_multi_head_attention_ops.h>
|
| 374 |
+
#include <ATen/ops/_triton_scaled_dot_attention_ops.h>
|
| 375 |
+
#include <ATen/ops/_unique_ops.h>
|
| 376 |
+
#include <ATen/ops/_unique2_ops.h>
|
| 377 |
+
#include <ATen/ops/_unpack_dual_ops.h>
|
| 378 |
+
#include <ATen/ops/_unsafe_index_ops.h>
|
| 379 |
+
#include <ATen/ops/_unsafe_index_put_ops.h>
|
| 380 |
+
#include <ATen/ops/_unsafe_masked_index_ops.h>
|
| 381 |
+
#include <ATen/ops/_unsafe_masked_index_put_accumulate_ops.h>
|
| 382 |
+
#include <ATen/ops/_unsafe_view_ops.h>
|
| 383 |
+
#include <ATen/ops/_upsample_bicubic2d_aa_ops.h>
|
| 384 |
+
#include <ATen/ops/_upsample_bicubic2d_aa_backward_ops.h>
|
| 385 |
+
#include <ATen/ops/_upsample_bilinear2d_aa_ops.h>
|
| 386 |
+
#include <ATen/ops/_upsample_bilinear2d_aa_backward_ops.h>
|
| 387 |
+
#include <ATen/ops/_upsample_nearest_exact1d_ops.h>
|
| 388 |
+
#include <ATen/ops/_upsample_nearest_exact1d_backward_ops.h>
|
| 389 |
+
#include <ATen/ops/_upsample_nearest_exact2d_ops.h>
|
| 390 |
+
#include <ATen/ops/_upsample_nearest_exact2d_backward_ops.h>
|
| 391 |
+
#include <ATen/ops/_upsample_nearest_exact3d_ops.h>
|
| 392 |
+
#include <ATen/ops/_upsample_nearest_exact3d_backward_ops.h>
|
| 393 |
+
#include <ATen/ops/_use_cudnn_ctc_loss_ops.h>
|
| 394 |
+
#include <ATen/ops/_use_cudnn_rnn_flatten_weight_ops.h>
|
| 395 |
+
#include <ATen/ops/_validate_compressed_sparse_indices_ops.h>
|
| 396 |
+
#include <ATen/ops/_validate_sparse_bsc_tensor_args_ops.h>
|
| 397 |
+
#include <ATen/ops/_validate_sparse_bsr_tensor_args_ops.h>
|
| 398 |
+
#include <ATen/ops/_validate_sparse_compressed_tensor_args_ops.h>
|
| 399 |
+
#include <ATen/ops/_validate_sparse_coo_tensor_args_ops.h>
|
| 400 |
+
#include <ATen/ops/_validate_sparse_csc_tensor_args_ops.h>
|
| 401 |
+
#include <ATen/ops/_validate_sparse_csr_tensor_args_ops.h>
|
| 402 |
+
#include <ATen/ops/_values_ops.h>
|
| 403 |
+
#include <ATen/ops/_values_copy_ops.h>
|
| 404 |
+
#include <ATen/ops/_version_ops.h>
|
| 405 |
+
#include <ATen/ops/_weight_int4pack_mm_ops.h>
|
| 406 |
+
#include <ATen/ops/_weight_int4pack_mm_for_cpu_ops.h>
|
| 407 |
+
#include <ATen/ops/_weight_int8pack_mm_ops.h>
|
| 408 |
+
#include <ATen/ops/_weight_norm_ops.h>
|
| 409 |
+
#include <ATen/ops/_weight_norm_differentiable_backward_ops.h>
|
| 410 |
+
#include <ATen/ops/_weight_norm_interface_ops.h>
|
| 411 |
+
#include <ATen/ops/_weight_norm_interface_backward_ops.h>
|
| 412 |
+
#include <ATen/ops/_wrapped_linear_prepack_ops.h>
|
| 413 |
+
#include <ATen/ops/_wrapped_quantized_linear_prepacked_ops.h>
|
| 414 |
+
#include <ATen/ops/abs_ops.h>
|
| 415 |
+
#include <ATen/ops/absolute_ops.h>
|
| 416 |
+
#include <ATen/ops/acos_ops.h>
|
| 417 |
+
#include <ATen/ops/acosh_ops.h>
|
| 418 |
+
#include <ATen/ops/adaptive_avg_pool1d_ops.h>
|
| 419 |
+
#include <ATen/ops/adaptive_avg_pool2d_ops.h>
|
| 420 |
+
#include <ATen/ops/adaptive_avg_pool3d_ops.h>
|
| 421 |
+
#include <ATen/ops/adaptive_avg_pool3d_backward_ops.h>
|
| 422 |
+
#include <ATen/ops/adaptive_max_pool1d_ops.h>
|
| 423 |
+
#include <ATen/ops/adaptive_max_pool2d_ops.h>
|
| 424 |
+
#include <ATen/ops/adaptive_max_pool2d_backward_ops.h>
|
| 425 |
+
#include <ATen/ops/adaptive_max_pool3d_ops.h>
|
| 426 |
+
#include <ATen/ops/adaptive_max_pool3d_backward_ops.h>
|
| 427 |
+
#include <ATen/ops/add_ops.h>
|
| 428 |
+
#include <ATen/ops/addbmm_ops.h>
|
| 429 |
+
#include <ATen/ops/addcdiv_ops.h>
|
| 430 |
+
#include <ATen/ops/addcmul_ops.h>
|
| 431 |
+
#include <ATen/ops/addmm_ops.h>
|
| 432 |
+
#include <ATen/ops/addmv_ops.h>
|
| 433 |
+
#include <ATen/ops/addr_ops.h>
|
| 434 |
+
#include <ATen/ops/adjoint_ops.h>
|
| 435 |
+
#include <ATen/ops/affine_grid_generator_ops.h>
|
| 436 |
+
#include <ATen/ops/affine_grid_generator_backward_ops.h>
|
| 437 |
+
#include <ATen/ops/alias_ops.h>
|
| 438 |
+
#include <ATen/ops/alias_copy_ops.h>
|
| 439 |
+
#include <ATen/ops/align_as_ops.h>
|
| 440 |
+
#include <ATen/ops/align_tensors_ops.h>
|
| 441 |
+
#include <ATen/ops/align_to_ops.h>
|
| 442 |
+
#include <ATen/ops/all_ops.h>
|
| 443 |
+
#include <ATen/ops/allclose_ops.h>
|
| 444 |
+
#include <ATen/ops/alpha_dropout_ops.h>
|
| 445 |
+
#include <ATen/ops/amax_ops.h>
|
| 446 |
+
#include <ATen/ops/amin_ops.h>
|
| 447 |
+
#include <ATen/ops/aminmax_ops.h>
|
| 448 |
+
#include <ATen/ops/and_ops.h>
|
| 449 |
+
#include <ATen/ops/angle_ops.h>
|
| 450 |
+
#include <ATen/ops/any_ops.h>
|
| 451 |
+
#include <ATen/ops/arange_ops.h>
|
| 452 |
+
#include <ATen/ops/arccos_ops.h>
|
| 453 |
+
#include <ATen/ops/arccosh_ops.h>
|
| 454 |
+
#include <ATen/ops/arcsin_ops.h>
|
| 455 |
+
#include <ATen/ops/arcsinh_ops.h>
|
| 456 |
+
#include <ATen/ops/arctan_ops.h>
|
| 457 |
+
#include <ATen/ops/arctan2_ops.h>
|
| 458 |
+
#include <ATen/ops/arctanh_ops.h>
|
| 459 |
+
#include <ATen/ops/argmax_ops.h>
|
| 460 |
+
#include <ATen/ops/argmin_ops.h>
|
| 461 |
+
#include <ATen/ops/argsort_ops.h>
|
| 462 |
+
#include <ATen/ops/argwhere_ops.h>
|
| 463 |
+
#include <ATen/ops/as_strided_ops.h>
|
| 464 |
+
#include <ATen/ops/as_strided_copy_ops.h>
|
| 465 |
+
#include <ATen/ops/as_strided_scatter_ops.h>
|
| 466 |
+
#include <ATen/ops/asin_ops.h>
|
| 467 |
+
#include <ATen/ops/asinh_ops.h>
|
| 468 |
+
#include <ATen/ops/atan_ops.h>
|
| 469 |
+
#include <ATen/ops/atan2_ops.h>
|
| 470 |
+
#include <ATen/ops/atanh_ops.h>
|
| 471 |
+
#include <ATen/ops/atleast_1d_ops.h>
|
| 472 |
+
#include <ATen/ops/atleast_2d_ops.h>
|
| 473 |
+
#include <ATen/ops/atleast_3d_ops.h>
|
| 474 |
+
#include <ATen/ops/avg_pool1d_ops.h>
|
| 475 |
+
#include <ATen/ops/avg_pool2d_ops.h>
|
| 476 |
+
#include <ATen/ops/avg_pool2d_backward_ops.h>
|
| 477 |
+
#include <ATen/ops/avg_pool3d_ops.h>
|
| 478 |
+
#include <ATen/ops/avg_pool3d_backward_ops.h>
|
| 479 |
+
#include <ATen/ops/baddbmm_ops.h>
|
| 480 |
+
#include <ATen/ops/bartlett_window_ops.h>
|
| 481 |
+
#include <ATen/ops/batch_norm_ops.h>
|
| 482 |
+
#include <ATen/ops/batch_norm_backward_ops.h>
|
| 483 |
+
#include <ATen/ops/batch_norm_backward_elemt_ops.h>
|
| 484 |
+
#include <ATen/ops/batch_norm_backward_reduce_ops.h>
|
| 485 |
+
#include <ATen/ops/batch_norm_elemt_ops.h>
|
| 486 |
+
#include <ATen/ops/batch_norm_gather_stats_ops.h>
|
| 487 |
+
#include <ATen/ops/batch_norm_gather_stats_with_counts_ops.h>
|
| 488 |
+
#include <ATen/ops/batch_norm_stats_ops.h>
|
| 489 |
+
#include <ATen/ops/batch_norm_update_stats_ops.h>
|
| 490 |
+
#include <ATen/ops/bernoulli_ops.h>
|
| 491 |
+
#include <ATen/ops/bilinear_ops.h>
|
| 492 |
+
#include <ATen/ops/binary_cross_entropy_ops.h>
|
| 493 |
+
#include <ATen/ops/binary_cross_entropy_backward_ops.h>
|
| 494 |
+
#include <ATen/ops/binary_cross_entropy_with_logits_ops.h>
|
| 495 |
+
#include <ATen/ops/bincount_ops.h>
|
| 496 |
+
#include <ATen/ops/binomial_ops.h>
|
| 497 |
+
#include <ATen/ops/bitwise_and_ops.h>
|
| 498 |
+
#include <ATen/ops/bitwise_left_shift_ops.h>
|
| 499 |
+
#include <ATen/ops/bitwise_not_ops.h>
|
| 500 |
+
#include <ATen/ops/bitwise_or_ops.h>
|
| 501 |
+
#include <ATen/ops/bitwise_right_shift_ops.h>
|
| 502 |
+
#include <ATen/ops/bitwise_xor_ops.h>
|
| 503 |
+
#include <ATen/ops/blackman_window_ops.h>
|
| 504 |
+
#include <ATen/ops/block_diag_ops.h>
|
| 505 |
+
#include <ATen/ops/bmm_ops.h>
|
| 506 |
+
#include <ATen/ops/broadcast_tensors_ops.h>
|
| 507 |
+
#include <ATen/ops/broadcast_to_ops.h>
|
| 508 |
+
#include <ATen/ops/bucketize_ops.h>
|
| 509 |
+
#include <ATen/ops/can_cast_ops.h>
|
| 510 |
+
#include <ATen/ops/cartesian_prod_ops.h>
|
| 511 |
+
#include <ATen/ops/cat_ops.h>
|
| 512 |
+
#include <ATen/ops/cauchy_ops.h>
|
| 513 |
+
#include <ATen/ops/ccol_indices_ops.h>
|
| 514 |
+
#include <ATen/ops/ccol_indices_copy_ops.h>
|
| 515 |
+
#include <ATen/ops/cdist_ops.h>
|
| 516 |
+
#include <ATen/ops/ceil_ops.h>
|
| 517 |
+
#include <ATen/ops/celu_ops.h>
|
| 518 |
+
#include <ATen/ops/chain_matmul_ops.h>
|
| 519 |
+
#include <ATen/ops/chalf_ops.h>
|
| 520 |
+
#include <ATen/ops/channel_shuffle_ops.h>
|
| 521 |
+
#include <ATen/ops/cholesky_ops.h>
|
| 522 |
+
#include <ATen/ops/cholesky_inverse_ops.h>
|
| 523 |
+
#include <ATen/ops/cholesky_solve_ops.h>
|
| 524 |
+
#include <ATen/ops/choose_qparams_optimized_ops.h>
|
| 525 |
+
#include <ATen/ops/chunk_ops.h>
|
| 526 |
+
#include <ATen/ops/clamp_ops.h>
|
| 527 |
+
#include <ATen/ops/clamp_max_ops.h>
|
| 528 |
+
#include <ATen/ops/clamp_min_ops.h>
|
| 529 |
+
#include <ATen/ops/clip_ops.h>
|
| 530 |
+
#include <ATen/ops/clone_ops.h>
|
| 531 |
+
#include <ATen/ops/coalesce_ops.h>
|
| 532 |
+
#include <ATen/ops/col2im_ops.h>
|
| 533 |
+
#include <ATen/ops/col_indices_ops.h>
|
| 534 |
+
#include <ATen/ops/col_indices_copy_ops.h>
|
| 535 |
+
#include <ATen/ops/column_stack_ops.h>
|
| 536 |
+
#include <ATen/ops/combinations_ops.h>
|
| 537 |
+
#include <ATen/ops/complex_ops.h>
|
| 538 |
+
#include <ATen/ops/concat_ops.h>
|
| 539 |
+
#include <ATen/ops/concatenate_ops.h>
|
| 540 |
+
#include <ATen/ops/conj_ops.h>
|
| 541 |
+
#include <ATen/ops/conj_physical_ops.h>
|
| 542 |
+
#include <ATen/ops/constant_pad_nd_ops.h>
|
| 543 |
+
#include <ATen/ops/contiguous_ops.h>
|
| 544 |
+
#include <ATen/ops/conv1d_ops.h>
|
| 545 |
+
#include <ATen/ops/conv2d_ops.h>
|
| 546 |
+
#include <ATen/ops/conv3d_ops.h>
|
| 547 |
+
#include <ATen/ops/conv_depthwise3d_ops.h>
|
| 548 |
+
#include <ATen/ops/conv_tbc_ops.h>
|
| 549 |
+
#include <ATen/ops/conv_tbc_backward_ops.h>
|
| 550 |
+
#include <ATen/ops/conv_transpose1d_ops.h>
|
| 551 |
+
#include <ATen/ops/conv_transpose2d_ops.h>
|
| 552 |
+
#include <ATen/ops/conv_transpose3d_ops.h>
|
| 553 |
+
#include <ATen/ops/convolution_ops.h>
|
| 554 |
+
#include <ATen/ops/convolution_backward_ops.h>
|
| 555 |
+
#include <ATen/ops/convolution_backward_overrideable_ops.h>
|
| 556 |
+
#include <ATen/ops/convolution_overrideable_ops.h>
|
| 557 |
+
#include <ATen/ops/copy_ops.h>
|
| 558 |
+
#include <ATen/ops/copy_sparse_to_sparse_ops.h>
|
| 559 |
+
#include <ATen/ops/copysign_ops.h>
|
| 560 |
+
#include <ATen/ops/corrcoef_ops.h>
|
| 561 |
+
#include <ATen/ops/cos_ops.h>
|
| 562 |
+
#include <ATen/ops/cosh_ops.h>
|
| 563 |
+
#include <ATen/ops/cosine_embedding_loss_ops.h>
|
| 564 |
+
#include <ATen/ops/cosine_similarity_ops.h>
|
| 565 |
+
#include <ATen/ops/count_nonzero_ops.h>
|
| 566 |
+
#include <ATen/ops/cov_ops.h>
|
| 567 |
+
#include <ATen/ops/cross_ops.h>
|
| 568 |
+
#include <ATen/ops/cross_entropy_loss_ops.h>
|
| 569 |
+
#include <ATen/ops/crow_indices_ops.h>
|
| 570 |
+
#include <ATen/ops/crow_indices_copy_ops.h>
|
| 571 |
+
#include <ATen/ops/ctc_loss_ops.h>
|
| 572 |
+
#include <ATen/ops/cudnn_affine_grid_generator_ops.h>
|
| 573 |
+
#include <ATen/ops/cudnn_affine_grid_generator_backward_ops.h>
|
| 574 |
+
#include <ATen/ops/cudnn_batch_norm_ops.h>
|
| 575 |
+
#include <ATen/ops/cudnn_batch_norm_backward_ops.h>
|
| 576 |
+
#include <ATen/ops/cudnn_convolution_ops.h>
|
| 577 |
+
#include <ATen/ops/cudnn_convolution_add_relu_ops.h>
|
| 578 |
+
#include <ATen/ops/cudnn_convolution_relu_ops.h>
|
| 579 |
+
#include <ATen/ops/cudnn_convolution_transpose_ops.h>
|
| 580 |
+
#include <ATen/ops/cudnn_grid_sampler_ops.h>
|
| 581 |
+
#include <ATen/ops/cudnn_grid_sampler_backward_ops.h>
|
| 582 |
+
#include <ATen/ops/cudnn_is_acceptable_ops.h>
|
| 583 |
+
#include <ATen/ops/cummax_ops.h>
|
| 584 |
+
#include <ATen/ops/cummaxmin_backward_ops.h>
|
| 585 |
+
#include <ATen/ops/cummin_ops.h>
|
| 586 |
+
#include <ATen/ops/cumprod_ops.h>
|
| 587 |
+
#include <ATen/ops/cumprod_backward_ops.h>
|
| 588 |
+
#include <ATen/ops/cumsum_ops.h>
|
| 589 |
+
#include <ATen/ops/cumulative_trapezoid_ops.h>
|
| 590 |
+
#include <ATen/ops/data_ops.h>
|
| 591 |
+
#include <ATen/ops/deg2rad_ops.h>
|
| 592 |
+
#include <ATen/ops/dense_dim_ops.h>
|
| 593 |
+
#include <ATen/ops/dequantize_ops.h>
|
| 594 |
+
#include <ATen/ops/det_ops.h>
|
| 595 |
+
#include <ATen/ops/detach_ops.h>
|
| 596 |
+
#include <ATen/ops/detach_copy_ops.h>
|
| 597 |
+
#include <ATen/ops/diag_ops.h>
|
| 598 |
+
#include <ATen/ops/diag_embed_ops.h>
|
| 599 |
+
#include <ATen/ops/diagflat_ops.h>
|
| 600 |
+
#include <ATen/ops/diagonal_ops.h>
|
| 601 |
+
#include <ATen/ops/diagonal_backward_ops.h>
|
| 602 |
+
#include <ATen/ops/diagonal_copy_ops.h>
|
| 603 |
+
#include <ATen/ops/diagonal_scatter_ops.h>
|
| 604 |
+
#include <ATen/ops/diff_ops.h>
|
| 605 |
+
#include <ATen/ops/digamma_ops.h>
|
| 606 |
+
#include <ATen/ops/dist_ops.h>
|
| 607 |
+
#include <ATen/ops/div_ops.h>
|
| 608 |
+
#include <ATen/ops/divide_ops.h>
|
| 609 |
+
#include <ATen/ops/dot_ops.h>
|
| 610 |
+
#include <ATen/ops/dropout_ops.h>
|
| 611 |
+
#include <ATen/ops/dsplit_ops.h>
|
| 612 |
+
#include <ATen/ops/dstack_ops.h>
|
| 613 |
+
#include <ATen/ops/einsum_ops.h>
|
| 614 |
+
#include <ATen/ops/elu_ops.h>
|
| 615 |
+
#include <ATen/ops/elu_backward_ops.h>
|
| 616 |
+
#include <ATen/ops/embedding_ops.h>
|
| 617 |
+
#include <ATen/ops/embedding_backward_ops.h>
|
| 618 |
+
#include <ATen/ops/embedding_bag_ops.h>
|
| 619 |
+
#include <ATen/ops/embedding_dense_backward_ops.h>
|
| 620 |
+
#include <ATen/ops/embedding_renorm_ops.h>
|
| 621 |
+
#include <ATen/ops/embedding_sparse_backward_ops.h>
|
| 622 |
+
#include <ATen/ops/empty_ops.h>
|
| 623 |
+
#include <ATen/ops/empty_like_ops.h>
|
| 624 |
+
#include <ATen/ops/empty_permuted_ops.h>
|
| 625 |
+
#include <ATen/ops/empty_quantized_ops.h>
|
| 626 |
+
#include <ATen/ops/empty_strided_ops.h>
|
| 627 |
+
#include <ATen/ops/eq_ops.h>
|
| 628 |
+
#include <ATen/ops/equal_ops.h>
|
| 629 |
+
#include <ATen/ops/erf_ops.h>
|
| 630 |
+
#include <ATen/ops/erfc_ops.h>
|
| 631 |
+
#include <ATen/ops/erfinv_ops.h>
|
| 632 |
+
#include <ATen/ops/exp_ops.h>
|
| 633 |
+
#include <ATen/ops/exp2_ops.h>
|
| 634 |
+
#include <ATen/ops/expand_ops.h>
|
| 635 |
+
#include <ATen/ops/expand_as_ops.h>
|
| 636 |
+
#include <ATen/ops/expand_copy_ops.h>
|
| 637 |
+
#include <ATen/ops/expm1_ops.h>
|
| 638 |
+
#include <ATen/ops/exponential_ops.h>
|
| 639 |
+
#include <ATen/ops/eye_ops.h>
|
| 640 |
+
#include <ATen/ops/fake_quantize_per_channel_affine_ops.h>
|
| 641 |
+
#include <ATen/ops/fake_quantize_per_channel_affine_cachemask_ops.h>
|
| 642 |
+
#include <ATen/ops/fake_quantize_per_channel_affine_cachemask_backward_ops.h>
|
| 643 |
+
#include <ATen/ops/fake_quantize_per_tensor_affine_ops.h>
|
| 644 |
+
#include <ATen/ops/fake_quantize_per_tensor_affine_cachemask_ops.h>
|
| 645 |
+
#include <ATen/ops/fake_quantize_per_tensor_affine_cachemask_backward_ops.h>
|
| 646 |
+
#include <ATen/ops/fbgemm_linear_fp16_weight_ops.h>
|
| 647 |
+
#include <ATen/ops/fbgemm_linear_fp16_weight_fp32_activation_ops.h>
|
| 648 |
+
#include <ATen/ops/fbgemm_linear_int8_weight_ops.h>
|
| 649 |
+
#include <ATen/ops/fbgemm_linear_int8_weight_fp32_activation_ops.h>
|
| 650 |
+
#include <ATen/ops/fbgemm_linear_quantize_weight_ops.h>
|
| 651 |
+
#include <ATen/ops/fbgemm_pack_gemm_matrix_fp16_ops.h>
|
| 652 |
+
#include <ATen/ops/fbgemm_pack_quantized_matrix_ops.h>
|
| 653 |
+
#include <ATen/ops/feature_alpha_dropout_ops.h>
|
| 654 |
+
#include <ATen/ops/feature_dropout_ops.h>
|
| 655 |
+
#include <ATen/ops/fft_fft_ops.h>
|
| 656 |
+
#include <ATen/ops/fft_fft2_ops.h>
|
| 657 |
+
#include <ATen/ops/fft_fftfreq_ops.h>
|
| 658 |
+
#include <ATen/ops/fft_fftn_ops.h>
|
| 659 |
+
#include <ATen/ops/fft_fftshift_ops.h>
|
| 660 |
+
#include <ATen/ops/fft_hfft_ops.h>
|
| 661 |
+
#include <ATen/ops/fft_hfft2_ops.h>
|
| 662 |
+
#include <ATen/ops/fft_hfftn_ops.h>
|
| 663 |
+
#include <ATen/ops/fft_ifft_ops.h>
|
| 664 |
+
#include <ATen/ops/fft_ifft2_ops.h>
|
| 665 |
+
#include <ATen/ops/fft_ifftn_ops.h>
|
| 666 |
+
#include <ATen/ops/fft_ifftshift_ops.h>
|
| 667 |
+
#include <ATen/ops/fft_ihfft_ops.h>
|
| 668 |
+
#include <ATen/ops/fft_ihfft2_ops.h>
|
| 669 |
+
#include <ATen/ops/fft_ihfftn_ops.h>
|
| 670 |
+
#include <ATen/ops/fft_irfft_ops.h>
|
| 671 |
+
#include <ATen/ops/fft_irfft2_ops.h>
|
| 672 |
+
#include <ATen/ops/fft_irfftn_ops.h>
|
| 673 |
+
#include <ATen/ops/fft_rfft_ops.h>
|
| 674 |
+
#include <ATen/ops/fft_rfft2_ops.h>
|
| 675 |
+
#include <ATen/ops/fft_rfftfreq_ops.h>
|
| 676 |
+
#include <ATen/ops/fft_rfftn_ops.h>
|
| 677 |
+
#include <ATen/ops/fill_ops.h>
|
| 678 |
+
#include <ATen/ops/fill_diagonal_ops.h>
|
| 679 |
+
#include <ATen/ops/fix_ops.h>
|
| 680 |
+
#include <ATen/ops/flatten_ops.h>
|
| 681 |
+
#include <ATen/ops/flatten_dense_tensors_ops.h>
|
| 682 |
+
#include <ATen/ops/flip_ops.h>
|
| 683 |
+
#include <ATen/ops/fliplr_ops.h>
|
| 684 |
+
#include <ATen/ops/flipud_ops.h>
|
| 685 |
+
#include <ATen/ops/float_power_ops.h>
|
| 686 |
+
#include <ATen/ops/floor_ops.h>
|
| 687 |
+
#include <ATen/ops/floor_divide_ops.h>
|
| 688 |
+
#include <ATen/ops/fmax_ops.h>
|
| 689 |
+
#include <ATen/ops/fmin_ops.h>
|
| 690 |
+
#include <ATen/ops/fmod_ops.h>
|
| 691 |
+
#include <ATen/ops/frac_ops.h>
|
| 692 |
+
#include <ATen/ops/fractional_max_pool2d_ops.h>
|
| 693 |
+
#include <ATen/ops/fractional_max_pool2d_backward_ops.h>
|
| 694 |
+
#include <ATen/ops/fractional_max_pool3d_ops.h>
|
| 695 |
+
#include <ATen/ops/fractional_max_pool3d_backward_ops.h>
|
| 696 |
+
#include <ATen/ops/frexp_ops.h>
|
| 697 |
+
#include <ATen/ops/frobenius_norm_ops.h>
|
| 698 |
+
#include <ATen/ops/from_file_ops.h>
|
| 699 |
+
#include <ATen/ops/full_ops.h>
|
| 700 |
+
#include <ATen/ops/full_like_ops.h>
|
| 701 |
+
#include <ATen/ops/fused_moving_avg_obs_fake_quant_ops.h>
|
| 702 |
+
#include <ATen/ops/gather_ops.h>
|
| 703 |
+
#include <ATen/ops/gather_backward_ops.h>
|
| 704 |
+
#include <ATen/ops/gcd_ops.h>
|
| 705 |
+
#include <ATen/ops/ge_ops.h>
|
| 706 |
+
#include <ATen/ops/gelu_ops.h>
|
| 707 |
+
#include <ATen/ops/gelu_backward_ops.h>
|
| 708 |
+
#include <ATen/ops/geometric_ops.h>
|
| 709 |
+
#include <ATen/ops/geqrf_ops.h>
|
| 710 |
+
#include <ATen/ops/ger_ops.h>
|
| 711 |
+
#include <ATen/ops/glu_ops.h>
|
| 712 |
+
#include <ATen/ops/glu_backward_ops.h>
|
| 713 |
+
#include <ATen/ops/glu_backward_jvp_ops.h>
|
| 714 |
+
#include <ATen/ops/glu_jvp_ops.h>
|
| 715 |
+
#include <ATen/ops/gradient_ops.h>
|
| 716 |
+
#include <ATen/ops/greater_ops.h>
|
| 717 |
+
#include <ATen/ops/greater_equal_ops.h>
|
| 718 |
+
#include <ATen/ops/grid_sampler_ops.h>
|
| 719 |
+
#include <ATen/ops/grid_sampler_2d_ops.h>
|
| 720 |
+
#include <ATen/ops/grid_sampler_2d_backward_ops.h>
|
| 721 |
+
#include <ATen/ops/grid_sampler_3d_ops.h>
|
| 722 |
+
#include <ATen/ops/grid_sampler_3d_backward_ops.h>
|
| 723 |
+
#include <ATen/ops/group_norm_ops.h>
|
| 724 |
+
#include <ATen/ops/gru_ops.h>
|
| 725 |
+
#include <ATen/ops/gru_cell_ops.h>
|
| 726 |
+
#include <ATen/ops/gt_ops.h>
|
| 727 |
+
#include <ATen/ops/hamming_window_ops.h>
|
| 728 |
+
#include <ATen/ops/hann_window_ops.h>
|
| 729 |
+
#include <ATen/ops/hardshrink_ops.h>
|
| 730 |
+
#include <ATen/ops/hardshrink_backward_ops.h>
|
| 731 |
+
#include <ATen/ops/hardsigmoid_ops.h>
|
| 732 |
+
#include <ATen/ops/hardsigmoid_backward_ops.h>
|
| 733 |
+
#include <ATen/ops/hardswish_ops.h>
|
| 734 |
+
#include <ATen/ops/hardswish_backward_ops.h>
|
| 735 |
+
#include <ATen/ops/hardtanh_ops.h>
|
| 736 |
+
#include <ATen/ops/hardtanh_backward_ops.h>
|
| 737 |
+
#include <ATen/ops/heaviside_ops.h>
|
| 738 |
+
#include <ATen/ops/hinge_embedding_loss_ops.h>
|
| 739 |
+
#include <ATen/ops/histc_ops.h>
|
| 740 |
+
#include <ATen/ops/histogram_ops.h>
|
| 741 |
+
#include <ATen/ops/histogramdd_ops.h>
|
| 742 |
+
#include <ATen/ops/hsplit_ops.h>
|
| 743 |
+
#include <ATen/ops/hspmm_ops.h>
|
| 744 |
+
#include <ATen/ops/hstack_ops.h>
|
| 745 |
+
#include <ATen/ops/huber_loss_ops.h>
|
| 746 |
+
#include <ATen/ops/huber_loss_backward_ops.h>
|
| 747 |
+
#include <ATen/ops/hypot_ops.h>
|
| 748 |
+
#include <ATen/ops/i0_ops.h>
|
| 749 |
+
#include <ATen/ops/igamma_ops.h>
|
| 750 |
+
#include <ATen/ops/igammac_ops.h>
|
| 751 |
+
#include <ATen/ops/im2col_ops.h>
|
| 752 |
+
#include <ATen/ops/imag_ops.h>
|
| 753 |
+
#include <ATen/ops/index_ops.h>
|
| 754 |
+
#include <ATen/ops/index_add_ops.h>
|
| 755 |
+
#include <ATen/ops/index_copy_ops.h>
|
| 756 |
+
#include <ATen/ops/index_fill_ops.h>
|
| 757 |
+
#include <ATen/ops/index_put_ops.h>
|
| 758 |
+
#include <ATen/ops/index_reduce_ops.h>
|
| 759 |
+
#include <ATen/ops/index_select_ops.h>
|
| 760 |
+
#include <ATen/ops/index_select_backward_ops.h>
|
| 761 |
+
#include <ATen/ops/indices_ops.h>
|
| 762 |
+
#include <ATen/ops/indices_copy_ops.h>
|
| 763 |
+
#include <ATen/ops/infinitely_differentiable_gelu_backward_ops.h>
|
| 764 |
+
#include <ATen/ops/inner_ops.h>
|
| 765 |
+
#include <ATen/ops/instance_norm_ops.h>
|
| 766 |
+
#include <ATen/ops/int_repr_ops.h>
|
| 767 |
+
#include <ATen/ops/inverse_ops.h>
|
| 768 |
+
#include <ATen/ops/is_coalesced_ops.h>
|
| 769 |
+
#include <ATen/ops/is_complex_ops.h>
|
| 770 |
+
#include <ATen/ops/is_conj_ops.h>
|
| 771 |
+
#include <ATen/ops/is_distributed_ops.h>
|
| 772 |
+
#include <ATen/ops/is_floating_point_ops.h>
|
| 773 |
+
#include <ATen/ops/is_inference_ops.h>
|
| 774 |
+
#include <ATen/ops/is_leaf_ops.h>
|
| 775 |
+
#include <ATen/ops/is_neg_ops.h>
|
| 776 |
+
#include <ATen/ops/is_nonzero_ops.h>
|
| 777 |
+
#include <ATen/ops/is_pinned_ops.h>
|
| 778 |
+
#include <ATen/ops/is_same_size_ops.h>
|
| 779 |
+
#include <ATen/ops/is_set_to_ops.h>
|
| 780 |
+
#include <ATen/ops/is_signed_ops.h>
|
| 781 |
+
#include <ATen/ops/is_vulkan_available_ops.h>
|
| 782 |
+
#include <ATen/ops/isclose_ops.h>
|
| 783 |
+
#include <ATen/ops/isfinite_ops.h>
|
| 784 |
+
#include <ATen/ops/isin_ops.h>
|
| 785 |
+
#include <ATen/ops/isinf_ops.h>
|
| 786 |
+
#include <ATen/ops/isnan_ops.h>
|
| 787 |
+
#include <ATen/ops/isneginf_ops.h>
|
| 788 |
+
#include <ATen/ops/isposinf_ops.h>
|
| 789 |
+
#include <ATen/ops/isreal_ops.h>
|
| 790 |
+
#include <ATen/ops/istft_ops.h>
|
| 791 |
+
#include <ATen/ops/item_ops.h>
|
| 792 |
+
#include <ATen/ops/kaiser_window_ops.h>
|
| 793 |
+
#include <ATen/ops/kl_div_ops.h>
|
| 794 |
+
#include <ATen/ops/kron_ops.h>
|
| 795 |
+
#include <ATen/ops/kthvalue_ops.h>
|
| 796 |
+
#include <ATen/ops/l1_loss_ops.h>
|
| 797 |
+
#include <ATen/ops/layer_norm_ops.h>
|
| 798 |
+
#include <ATen/ops/lcm_ops.h>
|
| 799 |
+
#include <ATen/ops/ldexp_ops.h>
|
| 800 |
+
#include <ATen/ops/le_ops.h>
|
| 801 |
+
#include <ATen/ops/leaky_relu_ops.h>
|
| 802 |
+
#include <ATen/ops/leaky_relu_backward_ops.h>
|
| 803 |
+
#include <ATen/ops/lerp_ops.h>
|
| 804 |
+
#include <ATen/ops/less_ops.h>
|
| 805 |
+
#include <ATen/ops/less_equal_ops.h>
|
| 806 |
+
#include <ATen/ops/lgamma_ops.h>
|
| 807 |
+
#include <ATen/ops/lift_ops.h>
|
| 808 |
+
#include <ATen/ops/lift_fresh_ops.h>
|
| 809 |
+
#include <ATen/ops/lift_fresh_copy_ops.h>
|
| 810 |
+
#include <ATen/ops/linalg_cholesky_ops.h>
|
| 811 |
+
#include <ATen/ops/linalg_cholesky_ex_ops.h>
|
| 812 |
+
#include <ATen/ops/linalg_cond_ops.h>
|
| 813 |
+
#include <ATen/ops/linalg_cross_ops.h>
|
| 814 |
+
#include <ATen/ops/linalg_det_ops.h>
|
| 815 |
+
#include <ATen/ops/linalg_diagonal_ops.h>
|
| 816 |
+
#include <ATen/ops/linalg_eig_ops.h>
|
| 817 |
+
#include <ATen/ops/linalg_eigh_ops.h>
|
| 818 |
+
#include <ATen/ops/linalg_eigvals_ops.h>
|
| 819 |
+
#include <ATen/ops/linalg_eigvalsh_ops.h>
|
| 820 |
+
#include <ATen/ops/linalg_householder_product_ops.h>
|
| 821 |
+
#include <ATen/ops/linalg_inv_ops.h>
|
| 822 |
+
#include <ATen/ops/linalg_inv_ex_ops.h>
|
| 823 |
+
#include <ATen/ops/linalg_ldl_factor_ops.h>
|
| 824 |
+
#include <ATen/ops/linalg_ldl_factor_ex_ops.h>
|
| 825 |
+
#include <ATen/ops/linalg_ldl_solve_ops.h>
|
| 826 |
+
#include <ATen/ops/linalg_lstsq_ops.h>
|
| 827 |
+
#include <ATen/ops/linalg_lu_ops.h>
|
| 828 |
+
#include <ATen/ops/linalg_lu_factor_ops.h>
|
| 829 |
+
#include <ATen/ops/linalg_lu_factor_ex_ops.h>
|
| 830 |
+
#include <ATen/ops/linalg_lu_solve_ops.h>
|
| 831 |
+
#include <ATen/ops/linalg_matmul_ops.h>
|
| 832 |
+
#include <ATen/ops/linalg_matrix_exp_ops.h>
|
| 833 |
+
#include <ATen/ops/linalg_matrix_norm_ops.h>
|
| 834 |
+
#include <ATen/ops/linalg_matrix_power_ops.h>
|
| 835 |
+
#include <ATen/ops/linalg_matrix_rank_ops.h>
|
| 836 |
+
#include <ATen/ops/linalg_multi_dot_ops.h>
|
| 837 |
+
#include <ATen/ops/linalg_norm_ops.h>
|
| 838 |
+
#include <ATen/ops/linalg_pinv_ops.h>
|
| 839 |
+
#include <ATen/ops/linalg_qr_ops.h>
|
| 840 |
+
#include <ATen/ops/linalg_slogdet_ops.h>
|
| 841 |
+
#include <ATen/ops/linalg_solve_ops.h>
|
| 842 |
+
#include <ATen/ops/linalg_solve_ex_ops.h>
|
| 843 |
+
#include <ATen/ops/linalg_solve_triangular_ops.h>
|
| 844 |
+
#include <ATen/ops/linalg_svd_ops.h>
|
| 845 |
+
#include <ATen/ops/linalg_svdvals_ops.h>
|
| 846 |
+
#include <ATen/ops/linalg_tensorinv_ops.h>
|
| 847 |
+
#include <ATen/ops/linalg_tensorsolve_ops.h>
|
| 848 |
+
#include <ATen/ops/linalg_vander_ops.h>
|
| 849 |
+
#include <ATen/ops/linalg_vecdot_ops.h>
|
| 850 |
+
#include <ATen/ops/linalg_vector_norm_ops.h>
|
| 851 |
+
#include <ATen/ops/linear_ops.h>
|
| 852 |
+
#include <ATen/ops/linear_backward_ops.h>
|
| 853 |
+
#include <ATen/ops/linspace_ops.h>
|
| 854 |
+
#include <ATen/ops/log_ops.h>
|
| 855 |
+
#include <ATen/ops/log10_ops.h>
|
| 856 |
+
#include <ATen/ops/log1p_ops.h>
|
| 857 |
+
#include <ATen/ops/log2_ops.h>
|
| 858 |
+
#include <ATen/ops/log_normal_ops.h>
|
| 859 |
+
#include <ATen/ops/log_sigmoid_ops.h>
|
| 860 |
+
#include <ATen/ops/log_sigmoid_backward_ops.h>
|
| 861 |
+
#include <ATen/ops/log_sigmoid_forward_ops.h>
|
| 862 |
+
#include <ATen/ops/log_softmax_ops.h>
|
| 863 |
+
#include <ATen/ops/logaddexp_ops.h>
|
| 864 |
+
#include <ATen/ops/logaddexp2_ops.h>
|
| 865 |
+
#include <ATen/ops/logcumsumexp_ops.h>
|
| 866 |
+
#include <ATen/ops/logdet_ops.h>
|
| 867 |
+
#include <ATen/ops/logical_and_ops.h>
|
| 868 |
+
#include <ATen/ops/logical_not_ops.h>
|
| 869 |
+
#include <ATen/ops/logical_or_ops.h>
|
| 870 |
+
#include <ATen/ops/logical_xor_ops.h>
|
| 871 |
+
#include <ATen/ops/logit_ops.h>
|
| 872 |
+
#include <ATen/ops/logit_backward_ops.h>
|
| 873 |
+
#include <ATen/ops/logspace_ops.h>
|
| 874 |
+
#include <ATen/ops/logsumexp_ops.h>
|
| 875 |
+
#include <ATen/ops/lshift_ops.h>
|
| 876 |
+
#include <ATen/ops/lstm_ops.h>
|
| 877 |
+
#include <ATen/ops/lstm_cell_ops.h>
|
| 878 |
+
#include <ATen/ops/lstm_mps_backward_ops.h>
|
| 879 |
+
#include <ATen/ops/lt_ops.h>
|
| 880 |
+
#include <ATen/ops/lu_solve_ops.h>
|
| 881 |
+
#include <ATen/ops/lu_unpack_ops.h>
|
| 882 |
+
#include <ATen/ops/mH_ops.h>
|
| 883 |
+
#include <ATen/ops/mT_ops.h>
|
| 884 |
+
#include <ATen/ops/margin_ranking_loss_ops.h>
|
| 885 |
+
#include <ATen/ops/masked_fill_ops.h>
|
| 886 |
+
#include <ATen/ops/masked_scatter_ops.h>
|
| 887 |
+
#include <ATen/ops/masked_scatter_backward_ops.h>
|
| 888 |
+
#include <ATen/ops/masked_select_ops.h>
|
| 889 |
+
#include <ATen/ops/masked_select_backward_ops.h>
|
| 890 |
+
#include <ATen/ops/matmul_ops.h>
|
| 891 |
+
#include <ATen/ops/matmul_backward_ops.h>
|
| 892 |
+
#include <ATen/ops/matrix_H_ops.h>
|
| 893 |
+
#include <ATen/ops/matrix_exp_ops.h>
|
| 894 |
+
#include <ATen/ops/matrix_exp_backward_ops.h>
|
| 895 |
+
#include <ATen/ops/matrix_power_ops.h>
|
| 896 |
+
#include <ATen/ops/max_ops.h>
|
| 897 |
+
#include <ATen/ops/max_pool1d_ops.h>
|
| 898 |
+
#include <ATen/ops/max_pool1d_with_indices_ops.h>
|
| 899 |
+
#include <ATen/ops/max_pool2d_ops.h>
|
| 900 |
+
#include <ATen/ops/max_pool2d_backward_ops.h>
|
| 901 |
+
#include <ATen/ops/max_pool2d_with_indices_ops.h>
|
| 902 |
+
#include <ATen/ops/max_pool2d_with_indices_backward_ops.h>
|
| 903 |
+
#include <ATen/ops/max_pool3d_ops.h>
|
| 904 |
+
#include <ATen/ops/max_pool3d_with_indices_ops.h>
|
| 905 |
+
#include <ATen/ops/max_pool3d_with_indices_backward_ops.h>
|
| 906 |
+
#include <ATen/ops/max_unpool2d_ops.h>
|
| 907 |
+
#include <ATen/ops/max_unpool3d_ops.h>
|
| 908 |
+
#include <ATen/ops/maximum_ops.h>
|
| 909 |
+
#include <ATen/ops/mean_ops.h>
|
| 910 |
+
#include <ATen/ops/median_ops.h>
|
| 911 |
+
#include <ATen/ops/meshgrid_ops.h>
|
| 912 |
+
#include <ATen/ops/min_ops.h>
|
| 913 |
+
#include <ATen/ops/minimum_ops.h>
|
| 914 |
+
#include <ATen/ops/miopen_batch_norm_ops.h>
|
| 915 |
+
#include <ATen/ops/miopen_batch_norm_backward_ops.h>
|
| 916 |
+
#include <ATen/ops/miopen_convolution_ops.h>
|
| 917 |
+
#include <ATen/ops/miopen_convolution_add_relu_ops.h>
|
| 918 |
+
#include <ATen/ops/miopen_convolution_relu_ops.h>
|
| 919 |
+
#include <ATen/ops/miopen_convolution_transpose_ops.h>
|
| 920 |
+
#include <ATen/ops/miopen_depthwise_convolution_ops.h>
|
| 921 |
+
#include <ATen/ops/miopen_rnn_ops.h>
|
| 922 |
+
#include <ATen/ops/miopen_rnn_backward_ops.h>
|
| 923 |
+
#include <ATen/ops/mish_ops.h>
|
| 924 |
+
#include <ATen/ops/mish_backward_ops.h>
|
| 925 |
+
#include <ATen/ops/mkldnn_adaptive_avg_pool2d_ops.h>
|
| 926 |
+
#include <ATen/ops/mkldnn_adaptive_avg_pool2d_backward_ops.h>
|
| 927 |
+
#include <ATen/ops/mkldnn_convolution_ops.h>
|
| 928 |
+
#include <ATen/ops/mkldnn_linear_ops.h>
|
| 929 |
+
#include <ATen/ops/mkldnn_linear_backward_ops.h>
|
| 930 |
+
#include <ATen/ops/mkldnn_linear_backward_input_ops.h>
|
| 931 |
+
#include <ATen/ops/mkldnn_linear_backward_weights_ops.h>
|
| 932 |
+
#include <ATen/ops/mkldnn_max_pool2d_ops.h>
|
| 933 |
+
#include <ATen/ops/mkldnn_max_pool2d_backward_ops.h>
|
| 934 |
+
#include <ATen/ops/mkldnn_max_pool3d_ops.h>
|
| 935 |
+
#include <ATen/ops/mkldnn_max_pool3d_backward_ops.h>
|
| 936 |
+
#include <ATen/ops/mkldnn_reorder_conv2d_weight_ops.h>
|
| 937 |
+
#include <ATen/ops/mkldnn_reorder_conv3d_weight_ops.h>
|
| 938 |
+
#include <ATen/ops/mkldnn_rnn_layer_ops.h>
|
| 939 |
+
#include <ATen/ops/mkldnn_rnn_layer_backward_ops.h>
|
| 940 |
+
#include <ATen/ops/mm_ops.h>
|
| 941 |
+
#include <ATen/ops/mode_ops.h>
|
| 942 |
+
#include <ATen/ops/moveaxis_ops.h>
|
| 943 |
+
#include <ATen/ops/movedim_ops.h>
|
| 944 |
+
#include <ATen/ops/mps_convolution_backward_ops.h>
|
| 945 |
+
#include <ATen/ops/mps_convolution_transpose_backward_ops.h>
|
| 946 |
+
#include <ATen/ops/mse_loss_ops.h>
|
| 947 |
+
#include <ATen/ops/mse_loss_backward_ops.h>
|
| 948 |
+
#include <ATen/ops/msort_ops.h>
|
| 949 |
+
#include <ATen/ops/mul_ops.h>
|
| 950 |
+
#include <ATen/ops/multi_margin_loss_ops.h>
|
| 951 |
+
#include <ATen/ops/multi_margin_loss_backward_ops.h>
|
| 952 |
+
#include <ATen/ops/multilabel_margin_loss_ops.h>
|
| 953 |
+
#include <ATen/ops/multilabel_margin_loss_backward_ops.h>
|
| 954 |
+
#include <ATen/ops/multilabel_margin_loss_forward_ops.h>
|
| 955 |
+
#include <ATen/ops/multinomial_ops.h>
|
| 956 |
+
#include <ATen/ops/multiply_ops.h>
|
| 957 |
+
#include <ATen/ops/mv_ops.h>
|
| 958 |
+
#include <ATen/ops/mvlgamma_ops.h>
|
| 959 |
+
#include <ATen/ops/nan_to_num_ops.h>
|
| 960 |
+
#include <ATen/ops/nanmean_ops.h>
|
| 961 |
+
#include <ATen/ops/nanmedian_ops.h>
|
| 962 |
+
#include <ATen/ops/nanquantile_ops.h>
|
| 963 |
+
#include <ATen/ops/nansum_ops.h>
|
| 964 |
+
#include <ATen/ops/narrow_ops.h>
|
| 965 |
+
#include <ATen/ops/narrow_copy_ops.h>
|
| 966 |
+
#include <ATen/ops/native_batch_norm_ops.h>
|
| 967 |
+
#include <ATen/ops/native_batch_norm_backward_ops.h>
|
| 968 |
+
#include <ATen/ops/native_channel_shuffle_ops.h>
|
| 969 |
+
#include <ATen/ops/native_dropout_ops.h>
|
| 970 |
+
#include <ATen/ops/native_dropout_backward_ops.h>
|
| 971 |
+
#include <ATen/ops/native_group_norm_ops.h>
|
| 972 |
+
#include <ATen/ops/native_group_norm_backward_ops.h>
|
| 973 |
+
#include <ATen/ops/native_layer_norm_ops.h>
|
| 974 |
+
#include <ATen/ops/native_layer_norm_backward_ops.h>
|
| 975 |
+
#include <ATen/ops/native_norm_ops.h>
|
| 976 |
+
#include <ATen/ops/ne_ops.h>
|
| 977 |
+
#include <ATen/ops/neg_ops.h>
|
| 978 |
+
#include <ATen/ops/negative_ops.h>
|
| 979 |
+
#include <ATen/ops/nested_to_padded_tensor_ops.h>
|
| 980 |
+
#include <ATen/ops/new_empty_ops.h>
|
| 981 |
+
#include <ATen/ops/new_empty_strided_ops.h>
|
| 982 |
+
#include <ATen/ops/new_full_ops.h>
|
| 983 |
+
#include <ATen/ops/new_ones_ops.h>
|
| 984 |
+
#include <ATen/ops/new_zeros_ops.h>
|
| 985 |
+
#include <ATen/ops/nextafter_ops.h>
|
| 986 |
+
#include <ATen/ops/nll_loss_ops.h>
|
| 987 |
+
#include <ATen/ops/nll_loss2d_ops.h>
|
| 988 |
+
#include <ATen/ops/nll_loss2d_backward_ops.h>
|
| 989 |
+
#include <ATen/ops/nll_loss2d_forward_ops.h>
|
| 990 |
+
#include <ATen/ops/nll_loss_backward_ops.h>
|
| 991 |
+
#include <ATen/ops/nll_loss_forward_ops.h>
|
| 992 |
+
#include <ATen/ops/nll_loss_nd_ops.h>
|
| 993 |
+
#include <ATen/ops/nonzero_ops.h>
|
| 994 |
+
#include <ATen/ops/nonzero_numpy_ops.h>
|
| 995 |
+
#include <ATen/ops/nonzero_static_ops.h>
|
| 996 |
+
#include <ATen/ops/norm_ops.h>
|
| 997 |
+
#include <ATen/ops/norm_except_dim_ops.h>
|
| 998 |
+
#include <ATen/ops/normal_ops.h>
|
| 999 |
+
#include <ATen/ops/not_equal_ops.h>
|
| 1000 |
+
#include <ATen/ops/nuclear_norm_ops.h>
|
| 1001 |
+
#include <ATen/ops/numpy_T_ops.h>
|
| 1002 |
+
#include <ATen/ops/one_hot_ops.h>
|
| 1003 |
+
#include <ATen/ops/ones_ops.h>
|
| 1004 |
+
#include <ATen/ops/ones_like_ops.h>
|
| 1005 |
+
#include <ATen/ops/or_ops.h>
|
| 1006 |
+
#include <ATen/ops/orgqr_ops.h>
|
| 1007 |
+
#include <ATen/ops/ormqr_ops.h>
|
| 1008 |
+
#include <ATen/ops/outer_ops.h>
|
| 1009 |
+
#include <ATen/ops/output_nr_ops.h>
|
| 1010 |
+
#include <ATen/ops/pad_ops.h>
|
| 1011 |
+
#include <ATen/ops/pad_sequence_ops.h>
|
| 1012 |
+
#include <ATen/ops/pairwise_distance_ops.h>
|
| 1013 |
+
#include <ATen/ops/pdist_ops.h>
|
| 1014 |
+
#include <ATen/ops/permute_ops.h>
|
| 1015 |
+
#include <ATen/ops/permute_copy_ops.h>
|
| 1016 |
+
#include <ATen/ops/pin_memory_ops.h>
|
| 1017 |
+
#include <ATen/ops/pinverse_ops.h>
|
| 1018 |
+
#include <ATen/ops/pixel_shuffle_ops.h>
|
| 1019 |
+
#include <ATen/ops/pixel_unshuffle_ops.h>
|
| 1020 |
+
#include <ATen/ops/poisson_ops.h>
|
| 1021 |
+
#include <ATen/ops/poisson_nll_loss_ops.h>
|
| 1022 |
+
#include <ATen/ops/polar_ops.h>
|
| 1023 |
+
#include <ATen/ops/polygamma_ops.h>
|
| 1024 |
+
#include <ATen/ops/positive_ops.h>
|
| 1025 |
+
#include <ATen/ops/pow_ops.h>
|
| 1026 |
+
#include <ATen/ops/prelu_ops.h>
|
| 1027 |
+
#include <ATen/ops/prod_ops.h>
|
| 1028 |
+
#include <ATen/ops/promote_types_ops.h>
|
| 1029 |
+
#include <ATen/ops/put_ops.h>
|
| 1030 |
+
#include <ATen/ops/q_per_channel_axis_ops.h>
|
| 1031 |
+
#include <ATen/ops/q_per_channel_scales_ops.h>
|
| 1032 |
+
#include <ATen/ops/q_per_channel_zero_points_ops.h>
|
| 1033 |
+
#include <ATen/ops/q_scale_ops.h>
|
| 1034 |
+
#include <ATen/ops/q_zero_point_ops.h>
|
| 1035 |
+
#include <ATen/ops/qr_ops.h>
|
| 1036 |
+
#include <ATen/ops/qscheme_ops.h>
|
| 1037 |
+
#include <ATen/ops/quantile_ops.h>
|
| 1038 |
+
#include <ATen/ops/quantize_per_channel_ops.h>
|
| 1039 |
+
#include <ATen/ops/quantize_per_tensor_ops.h>
|
| 1040 |
+
#include <ATen/ops/quantize_per_tensor_dynamic_ops.h>
|
| 1041 |
+
#include <ATen/ops/quantized_batch_norm_ops.h>
|
| 1042 |
+
#include <ATen/ops/quantized_gru_cell_ops.h>
|
| 1043 |
+
#include <ATen/ops/quantized_lstm_cell_ops.h>
|
| 1044 |
+
#include <ATen/ops/quantized_max_pool1d_ops.h>
|
| 1045 |
+
#include <ATen/ops/quantized_max_pool2d_ops.h>
|
| 1046 |
+
#include <ATen/ops/quantized_max_pool3d_ops.h>
|
| 1047 |
+
#include <ATen/ops/quantized_rnn_relu_cell_ops.h>
|
| 1048 |
+
#include <ATen/ops/quantized_rnn_tanh_cell_ops.h>
|
| 1049 |
+
#include <ATen/ops/rad2deg_ops.h>
|
| 1050 |
+
#include <ATen/ops/rand_ops.h>
|
| 1051 |
+
#include <ATen/ops/rand_like_ops.h>
|
| 1052 |
+
#include <ATen/ops/randint_ops.h>
|
| 1053 |
+
#include <ATen/ops/randint_like_ops.h>
|
| 1054 |
+
#include <ATen/ops/randn_ops.h>
|
| 1055 |
+
#include <ATen/ops/randn_like_ops.h>
|
| 1056 |
+
#include <ATen/ops/random_ops.h>
|
| 1057 |
+
#include <ATen/ops/randperm_ops.h>
|
| 1058 |
+
#include <ATen/ops/range_ops.h>
|
| 1059 |
+
#include <ATen/ops/ravel_ops.h>
|
| 1060 |
+
#include <ATen/ops/real_ops.h>
|
| 1061 |
+
#include <ATen/ops/reciprocal_ops.h>
|
| 1062 |
+
#include <ATen/ops/record_stream_ops.h>
|
| 1063 |
+
#include <ATen/ops/refine_names_ops.h>
|
| 1064 |
+
#include <ATen/ops/reflection_pad1d_ops.h>
|
| 1065 |
+
#include <ATen/ops/reflection_pad1d_backward_ops.h>
|
| 1066 |
+
#include <ATen/ops/reflection_pad2d_ops.h>
|
| 1067 |
+
#include <ATen/ops/reflection_pad2d_backward_ops.h>
|
| 1068 |
+
#include <ATen/ops/reflection_pad3d_ops.h>
|
| 1069 |
+
#include <ATen/ops/reflection_pad3d_backward_ops.h>
|
| 1070 |
+
#include <ATen/ops/relu_ops.h>
|
| 1071 |
+
#include <ATen/ops/relu6_ops.h>
|
| 1072 |
+
#include <ATen/ops/remainder_ops.h>
|
| 1073 |
+
#include <ATen/ops/rename_ops.h>
|
| 1074 |
+
#include <ATen/ops/renorm_ops.h>
|
| 1075 |
+
#include <ATen/ops/repeat_ops.h>
|
| 1076 |
+
#include <ATen/ops/repeat_interleave_ops.h>
|
| 1077 |
+
#include <ATen/ops/replication_pad1d_ops.h>
|
| 1078 |
+
#include <ATen/ops/replication_pad1d_backward_ops.h>
|
| 1079 |
+
#include <ATen/ops/replication_pad2d_ops.h>
|
| 1080 |
+
#include <ATen/ops/replication_pad2d_backward_ops.h>
|
| 1081 |
+
#include <ATen/ops/replication_pad3d_ops.h>
|
| 1082 |
+
#include <ATen/ops/replication_pad3d_backward_ops.h>
|
| 1083 |
+
#include <ATen/ops/requires_grad_ops.h>
|
| 1084 |
+
#include <ATen/ops/reshape_ops.h>
|
| 1085 |
+
#include <ATen/ops/reshape_as_ops.h>
|
| 1086 |
+
#include <ATen/ops/resize_ops.h>
|
| 1087 |
+
#include <ATen/ops/resize_as_ops.h>
|
| 1088 |
+
#include <ATen/ops/resize_as_sparse_ops.h>
|
| 1089 |
+
#include <ATen/ops/resolve_conj_ops.h>
|
| 1090 |
+
#include <ATen/ops/resolve_neg_ops.h>
|
| 1091 |
+
#include <ATen/ops/result_type_ops.h>
|
| 1092 |
+
#include <ATen/ops/retain_grad_ops.h>
|
| 1093 |
+
#include <ATen/ops/retains_grad_ops.h>
|
| 1094 |
+
#include <ATen/ops/rms_norm_ops.h>
|
| 1095 |
+
#include <ATen/ops/rnn_relu_ops.h>
|
| 1096 |
+
#include <ATen/ops/rnn_relu_cell_ops.h>
|
| 1097 |
+
#include <ATen/ops/rnn_tanh_ops.h>
|
| 1098 |
+
#include <ATen/ops/rnn_tanh_cell_ops.h>
|
| 1099 |
+
#include <ATen/ops/roll_ops.h>
|
| 1100 |
+
#include <ATen/ops/rot90_ops.h>
|
| 1101 |
+
#include <ATen/ops/round_ops.h>
|
| 1102 |
+
#include <ATen/ops/row_indices_ops.h>
|
| 1103 |
+
#include <ATen/ops/row_indices_copy_ops.h>
|
| 1104 |
+
#include <ATen/ops/row_stack_ops.h>
|
| 1105 |
+
#include <ATen/ops/rrelu_ops.h>
|
| 1106 |
+
#include <ATen/ops/rrelu_with_noise_ops.h>
|
| 1107 |
+
#include <ATen/ops/rrelu_with_noise_backward_ops.h>
|
| 1108 |
+
#include <ATen/ops/rshift_ops.h>
|
| 1109 |
+
#include <ATen/ops/rsqrt_ops.h>
|
| 1110 |
+
#include <ATen/ops/rsub_ops.h>
|
| 1111 |
+
#include <ATen/ops/scalar_tensor_ops.h>
|
| 1112 |
+
#include <ATen/ops/scaled_dot_product_attention_ops.h>
|
| 1113 |
+
#include <ATen/ops/scatter_ops.h>
|
| 1114 |
+
#include <ATen/ops/scatter_add_ops.h>
|
| 1115 |
+
#include <ATen/ops/scatter_reduce_ops.h>
|
| 1116 |
+
#include <ATen/ops/searchsorted_ops.h>
|
| 1117 |
+
#include <ATen/ops/segment_reduce_ops.h>
|
| 1118 |
+
#include <ATen/ops/select_ops.h>
|
| 1119 |
+
#include <ATen/ops/select_backward_ops.h>
|
| 1120 |
+
#include <ATen/ops/select_copy_ops.h>
|
| 1121 |
+
#include <ATen/ops/select_scatter_ops.h>
|
| 1122 |
+
#include <ATen/ops/selu_ops.h>
|
| 1123 |
+
#include <ATen/ops/set_ops.h>
|
| 1124 |
+
#include <ATen/ops/set_data_ops.h>
|
| 1125 |
+
#include <ATen/ops/sgn_ops.h>
|
| 1126 |
+
#include <ATen/ops/sigmoid_ops.h>
|
| 1127 |
+
#include <ATen/ops/sigmoid_backward_ops.h>
|
| 1128 |
+
#include <ATen/ops/sign_ops.h>
|
| 1129 |
+
#include <ATen/ops/signbit_ops.h>
|
| 1130 |
+
#include <ATen/ops/silu_ops.h>
|
| 1131 |
+
#include <ATen/ops/silu_backward_ops.h>
|
| 1132 |
+
#include <ATen/ops/sin_ops.h>
|
| 1133 |
+
#include <ATen/ops/sinc_ops.h>
|
| 1134 |
+
#include <ATen/ops/sinh_ops.h>
|
| 1135 |
+
#include <ATen/ops/size_ops.h>
|
| 1136 |
+
#include <ATen/ops/slice_ops.h>
|
| 1137 |
+
#include <ATen/ops/slice_backward_ops.h>
|
| 1138 |
+
#include <ATen/ops/slice_copy_ops.h>
|
| 1139 |
+
#include <ATen/ops/slice_inverse_ops.h>
|
| 1140 |
+
#include <ATen/ops/slice_scatter_ops.h>
|
| 1141 |
+
#include <ATen/ops/slogdet_ops.h>
|
| 1142 |
+
#include <ATen/ops/slow_conv3d_ops.h>
|
| 1143 |
+
#include <ATen/ops/slow_conv3d_forward_ops.h>
|
| 1144 |
+
#include <ATen/ops/slow_conv_dilated2d_ops.h>
|
| 1145 |
+
#include <ATen/ops/slow_conv_dilated3d_ops.h>
|
| 1146 |
+
#include <ATen/ops/slow_conv_transpose2d_ops.h>
|
| 1147 |
+
#include <ATen/ops/slow_conv_transpose3d_ops.h>
|
| 1148 |
+
#include <ATen/ops/smm_ops.h>
|
| 1149 |
+
#include <ATen/ops/smooth_l1_loss_ops.h>
|
| 1150 |
+
#include <ATen/ops/smooth_l1_loss_backward_ops.h>
|
| 1151 |
+
#include <ATen/ops/soft_margin_loss_ops.h>
|
| 1152 |
+
#include <ATen/ops/soft_margin_loss_backward_ops.h>
|
| 1153 |
+
#include <ATen/ops/softmax_ops.h>
|
| 1154 |
+
#include <ATen/ops/softplus_ops.h>
|
| 1155 |
+
#include <ATen/ops/softplus_backward_ops.h>
|
| 1156 |
+
#include <ATen/ops/softshrink_ops.h>
|
| 1157 |
+
#include <ATen/ops/softshrink_backward_ops.h>
|
| 1158 |
+
#include <ATen/ops/sort_ops.h>
|
| 1159 |
+
#include <ATen/ops/sparse_bsc_tensor_ops.h>
|
| 1160 |
+
#include <ATen/ops/sparse_bsr_tensor_ops.h>
|
| 1161 |
+
#include <ATen/ops/sparse_compressed_tensor_ops.h>
|
| 1162 |
+
#include <ATen/ops/sparse_coo_tensor_ops.h>
|
| 1163 |
+
#include <ATen/ops/sparse_csc_tensor_ops.h>
|
| 1164 |
+
#include <ATen/ops/sparse_csr_tensor_ops.h>
|
| 1165 |
+
#include <ATen/ops/sparse_dim_ops.h>
|
| 1166 |
+
#include <ATen/ops/sparse_mask_ops.h>
|
| 1167 |
+
#include <ATen/ops/sparse_resize_ops.h>
|
| 1168 |
+
#include <ATen/ops/sparse_resize_and_clear_ops.h>
|
| 1169 |
+
#include <ATen/ops/sparse_sampled_addmm_ops.h>
|
| 1170 |
+
#include <ATen/ops/special_airy_ai_ops.h>
|
| 1171 |
+
#include <ATen/ops/special_bessel_j0_ops.h>
|
| 1172 |
+
#include <ATen/ops/special_bessel_j1_ops.h>
|
| 1173 |
+
#include <ATen/ops/special_bessel_y0_ops.h>
|
| 1174 |
+
#include <ATen/ops/special_bessel_y1_ops.h>
|
| 1175 |
+
#include <ATen/ops/special_chebyshev_polynomial_t_ops.h>
|
| 1176 |
+
#include <ATen/ops/special_chebyshev_polynomial_u_ops.h>
|
| 1177 |
+
#include <ATen/ops/special_chebyshev_polynomial_v_ops.h>
|
| 1178 |
+
#include <ATen/ops/special_chebyshev_polynomial_w_ops.h>
|
| 1179 |
+
#include <ATen/ops/special_digamma_ops.h>
|
| 1180 |
+
#include <ATen/ops/special_entr_ops.h>
|
| 1181 |
+
#include <ATen/ops/special_erf_ops.h>
|
| 1182 |
+
#include <ATen/ops/special_erfc_ops.h>
|
| 1183 |
+
#include <ATen/ops/special_erfcx_ops.h>
|
| 1184 |
+
#include <ATen/ops/special_erfinv_ops.h>
|
| 1185 |
+
#include <ATen/ops/special_exp2_ops.h>
|
| 1186 |
+
#include <ATen/ops/special_expit_ops.h>
|
| 1187 |
+
#include <ATen/ops/special_expm1_ops.h>
|
| 1188 |
+
#include <ATen/ops/special_gammainc_ops.h>
|
| 1189 |
+
#include <ATen/ops/special_gammaincc_ops.h>
|
| 1190 |
+
#include <ATen/ops/special_gammaln_ops.h>
|
| 1191 |
+
#include <ATen/ops/special_hermite_polynomial_h_ops.h>
|
| 1192 |
+
#include <ATen/ops/special_hermite_polynomial_he_ops.h>
|
| 1193 |
+
#include <ATen/ops/special_i0_ops.h>
|
| 1194 |
+
#include <ATen/ops/special_i0e_ops.h>
|
| 1195 |
+
#include <ATen/ops/special_i1_ops.h>
|
| 1196 |
+
#include <ATen/ops/special_i1e_ops.h>
|
| 1197 |
+
#include <ATen/ops/special_laguerre_polynomial_l_ops.h>
|
| 1198 |
+
#include <ATen/ops/special_legendre_polynomial_p_ops.h>
|
| 1199 |
+
#include <ATen/ops/special_log1p_ops.h>
|
| 1200 |
+
#include <ATen/ops/special_log_ndtr_ops.h>
|
| 1201 |
+
#include <ATen/ops/special_log_softmax_ops.h>
|
| 1202 |
+
#include <ATen/ops/special_logit_ops.h>
|
| 1203 |
+
#include <ATen/ops/special_logsumexp_ops.h>
|
| 1204 |
+
#include <ATen/ops/special_modified_bessel_i0_ops.h>
|
| 1205 |
+
#include <ATen/ops/special_modified_bessel_i1_ops.h>
|
| 1206 |
+
#include <ATen/ops/special_modified_bessel_k0_ops.h>
|
| 1207 |
+
#include <ATen/ops/special_modified_bessel_k1_ops.h>
|
| 1208 |
+
#include <ATen/ops/special_multigammaln_ops.h>
|
| 1209 |
+
#include <ATen/ops/special_ndtr_ops.h>
|
| 1210 |
+
#include <ATen/ops/special_ndtri_ops.h>
|
| 1211 |
+
#include <ATen/ops/special_polygamma_ops.h>
|
| 1212 |
+
#include <ATen/ops/special_psi_ops.h>
|
| 1213 |
+
#include <ATen/ops/special_round_ops.h>
|
| 1214 |
+
#include <ATen/ops/special_scaled_modified_bessel_k0_ops.h>
|
| 1215 |
+
#include <ATen/ops/special_scaled_modified_bessel_k1_ops.h>
|
| 1216 |
+
#include <ATen/ops/special_shifted_chebyshev_polynomial_t_ops.h>
|
| 1217 |
+
#include <ATen/ops/special_shifted_chebyshev_polynomial_u_ops.h>
|
| 1218 |
+
#include <ATen/ops/special_shifted_chebyshev_polynomial_v_ops.h>
|
| 1219 |
+
#include <ATen/ops/special_shifted_chebyshev_polynomial_w_ops.h>
|
| 1220 |
+
#include <ATen/ops/special_sinc_ops.h>
|
| 1221 |
+
#include <ATen/ops/special_softmax_ops.h>
|
| 1222 |
+
#include <ATen/ops/special_spherical_bessel_j0_ops.h>
|
| 1223 |
+
#include <ATen/ops/special_xlog1py_ops.h>
|
| 1224 |
+
#include <ATen/ops/special_xlogy_ops.h>
|
| 1225 |
+
#include <ATen/ops/special_zeta_ops.h>
|
| 1226 |
+
#include <ATen/ops/split_ops.h>
|
| 1227 |
+
#include <ATen/ops/split_copy_ops.h>
|
| 1228 |
+
#include <ATen/ops/split_with_sizes_ops.h>
|
| 1229 |
+
#include <ATen/ops/split_with_sizes_copy_ops.h>
|
| 1230 |
+
#include <ATen/ops/sqrt_ops.h>
|
| 1231 |
+
#include <ATen/ops/square_ops.h>
|
| 1232 |
+
#include <ATen/ops/squeeze_ops.h>
|
| 1233 |
+
#include <ATen/ops/squeeze_copy_ops.h>
|
| 1234 |
+
#include <ATen/ops/sspaddmm_ops.h>
|
| 1235 |
+
#include <ATen/ops/stack_ops.h>
|
| 1236 |
+
#include <ATen/ops/std_ops.h>
|
| 1237 |
+
#include <ATen/ops/std_mean_ops.h>
|
| 1238 |
+
#include <ATen/ops/stft_ops.h>
|
| 1239 |
+
#include <ATen/ops/stride_ops.h>
|
| 1240 |
+
#include <ATen/ops/sub_ops.h>
|
| 1241 |
+
#include <ATen/ops/subtract_ops.h>
|
| 1242 |
+
#include <ATen/ops/sum_ops.h>
|
| 1243 |
+
#include <ATen/ops/sum_to_size_ops.h>
|
| 1244 |
+
#include <ATen/ops/svd_ops.h>
|
| 1245 |
+
#include <ATen/ops/swapaxes_ops.h>
|
| 1246 |
+
#include <ATen/ops/swapdims_ops.h>
|
| 1247 |
+
#include <ATen/ops/sym_constrain_range_ops.h>
|
| 1248 |
+
#include <ATen/ops/sym_constrain_range_for_size_ops.h>
|
| 1249 |
+
#include <ATen/ops/sym_numel_ops.h>
|
| 1250 |
+
#include <ATen/ops/sym_size_ops.h>
|
| 1251 |
+
#include <ATen/ops/sym_storage_offset_ops.h>
|
| 1252 |
+
#include <ATen/ops/sym_stride_ops.h>
|
| 1253 |
+
#include <ATen/ops/t_ops.h>
|
| 1254 |
+
#include <ATen/ops/t_copy_ops.h>
|
| 1255 |
+
#include <ATen/ops/take_ops.h>
|
| 1256 |
+
#include <ATen/ops/take_along_dim_ops.h>
|
| 1257 |
+
#include <ATen/ops/tan_ops.h>
|
| 1258 |
+
#include <ATen/ops/tanh_ops.h>
|
| 1259 |
+
#include <ATen/ops/tanh_backward_ops.h>
|
| 1260 |
+
#include <ATen/ops/tensor_split_ops.h>
|
| 1261 |
+
#include <ATen/ops/tensordot_ops.h>
|
| 1262 |
+
#include <ATen/ops/thnn_conv2d_ops.h>
|
| 1263 |
+
#include <ATen/ops/threshold_ops.h>
|
| 1264 |
+
#include <ATen/ops/threshold_backward_ops.h>
|
| 1265 |
+
#include <ATen/ops/tile_ops.h>
|
| 1266 |
+
#include <ATen/ops/to_ops.h>
|
| 1267 |
+
#include <ATen/ops/to_dense_ops.h>
|
| 1268 |
+
#include <ATen/ops/to_dense_backward_ops.h>
|
| 1269 |
+
#include <ATen/ops/to_mkldnn_ops.h>
|
| 1270 |
+
#include <ATen/ops/to_mkldnn_backward_ops.h>
|
| 1271 |
+
#include <ATen/ops/to_padded_tensor_ops.h>
|
| 1272 |
+
#include <ATen/ops/to_sparse_ops.h>
|
| 1273 |
+
#include <ATen/ops/to_sparse_bsc_ops.h>
|
| 1274 |
+
#include <ATen/ops/to_sparse_bsr_ops.h>
|
| 1275 |
+
#include <ATen/ops/to_sparse_csc_ops.h>
|
| 1276 |
+
#include <ATen/ops/to_sparse_csr_ops.h>
|
| 1277 |
+
#include <ATen/ops/topk_ops.h>
|
| 1278 |
+
#include <ATen/ops/trace_ops.h>
|
| 1279 |
+
#include <ATen/ops/trace_backward_ops.h>
|
| 1280 |
+
#include <ATen/ops/transpose_ops.h>
|
| 1281 |
+
#include <ATen/ops/transpose_copy_ops.h>
|
| 1282 |
+
#include <ATen/ops/trapezoid_ops.h>
|
| 1283 |
+
#include <ATen/ops/trapz_ops.h>
|
| 1284 |
+
#include <ATen/ops/triangular_solve_ops.h>
|
| 1285 |
+
#include <ATen/ops/tril_ops.h>
|
| 1286 |
+
#include <ATen/ops/tril_indices_ops.h>
|
| 1287 |
+
#include <ATen/ops/triplet_margin_loss_ops.h>
|
| 1288 |
+
#include <ATen/ops/triu_ops.h>
|
| 1289 |
+
#include <ATen/ops/triu_indices_ops.h>
|
| 1290 |
+
#include <ATen/ops/true_divide_ops.h>
|
| 1291 |
+
#include <ATen/ops/trunc_ops.h>
|
| 1292 |
+
#include <ATen/ops/type_as_ops.h>
|
| 1293 |
+
#include <ATen/ops/unbind_ops.h>
|
| 1294 |
+
#include <ATen/ops/unbind_copy_ops.h>
|
| 1295 |
+
#include <ATen/ops/unflatten_ops.h>
|
| 1296 |
+
#include <ATen/ops/unflatten_dense_tensors_ops.h>
|
| 1297 |
+
#include <ATen/ops/unfold_ops.h>
|
| 1298 |
+
#include <ATen/ops/unfold_backward_ops.h>
|
| 1299 |
+
#include <ATen/ops/unfold_copy_ops.h>
|
| 1300 |
+
#include <ATen/ops/uniform_ops.h>
|
| 1301 |
+
#include <ATen/ops/unique_consecutive_ops.h>
|
| 1302 |
+
#include <ATen/ops/unique_dim_ops.h>
|
| 1303 |
+
#include <ATen/ops/unique_dim_consecutive_ops.h>
|
| 1304 |
+
#include <ATen/ops/unsafe_chunk_ops.h>
|
| 1305 |
+
#include <ATen/ops/unsafe_split_ops.h>
|
| 1306 |
+
#include <ATen/ops/unsafe_split_with_sizes_ops.h>
|
| 1307 |
+
#include <ATen/ops/unsqueeze_ops.h>
|
| 1308 |
+
#include <ATen/ops/unsqueeze_copy_ops.h>
|
| 1309 |
+
#include <ATen/ops/upsample_bicubic2d_ops.h>
|
| 1310 |
+
#include <ATen/ops/upsample_bicubic2d_backward_ops.h>
|
| 1311 |
+
#include <ATen/ops/upsample_bilinear2d_ops.h>
|
| 1312 |
+
#include <ATen/ops/upsample_bilinear2d_backward_ops.h>
|
| 1313 |
+
#include <ATen/ops/upsample_linear1d_ops.h>
|
| 1314 |
+
#include <ATen/ops/upsample_linear1d_backward_ops.h>
|
| 1315 |
+
#include <ATen/ops/upsample_nearest1d_ops.h>
|
| 1316 |
+
#include <ATen/ops/upsample_nearest1d_backward_ops.h>
|
| 1317 |
+
#include <ATen/ops/upsample_nearest2d_ops.h>
|
| 1318 |
+
#include <ATen/ops/upsample_nearest2d_backward_ops.h>
|
| 1319 |
+
#include <ATen/ops/upsample_nearest3d_ops.h>
|
| 1320 |
+
#include <ATen/ops/upsample_nearest3d_backward_ops.h>
|
| 1321 |
+
#include <ATen/ops/upsample_trilinear3d_ops.h>
|
| 1322 |
+
#include <ATen/ops/upsample_trilinear3d_backward_ops.h>
|
| 1323 |
+
#include <ATen/ops/value_selecting_reduction_backward_ops.h>
|
| 1324 |
+
#include <ATen/ops/values_ops.h>
|
| 1325 |
+
#include <ATen/ops/values_copy_ops.h>
|
| 1326 |
+
#include <ATen/ops/vander_ops.h>
|
| 1327 |
+
#include <ATen/ops/var_ops.h>
|
| 1328 |
+
#include <ATen/ops/var_mean_ops.h>
|
| 1329 |
+
#include <ATen/ops/vdot_ops.h>
|
| 1330 |
+
#include <ATen/ops/view_ops.h>
|
| 1331 |
+
#include <ATen/ops/view_as_ops.h>
|
| 1332 |
+
#include <ATen/ops/view_as_complex_ops.h>
|
| 1333 |
+
#include <ATen/ops/view_as_complex_copy_ops.h>
|
| 1334 |
+
#include <ATen/ops/view_as_real_ops.h>
|
| 1335 |
+
#include <ATen/ops/view_as_real_copy_ops.h>
|
| 1336 |
+
#include <ATen/ops/view_copy_ops.h>
|
| 1337 |
+
#include <ATen/ops/vsplit_ops.h>
|
| 1338 |
+
#include <ATen/ops/vstack_ops.h>
|
| 1339 |
+
#include <ATen/ops/where_ops.h>
|
| 1340 |
+
#include <ATen/ops/xlogy_ops.h>
|
| 1341 |
+
#include <ATen/ops/xor_ops.h>
|
| 1342 |
+
#include <ATen/ops/zero_ops.h>
|
| 1343 |
+
#include <ATen/ops/zeros_ops.h>
|
| 1344 |
+
#include <ATen/ops/zeros_like_ops.h>
|
| 1345 |
+
|
| 1346 |
+
// Extension writers: do you write wrapper functions? Are you frustrated with
|
| 1347 |
+
// resolving overloads of operators? Are you frustrated with dealing with
|
| 1348 |
+
// pointer-to-methods and resolving overloads of pointer-to-methods?? Look no
|
| 1349 |
+
// further, this is the utility for you.
|
| 1350 |
+
//
|
| 1351 |
+
// Given an operator schema: aten::op.overload(...
|
| 1352 |
+
//
|
| 1353 |
+
// Use ATEN_FN2(op, overload) to get a *function* version of the operator
|
| 1354 |
+
// that is guaranteed to not be overloaded. This means that you can safely
|
| 1355 |
+
// decltype(&ATEN_FN2(op, overload)) it. NB: the 2 means this macro takes 2 args.
|
| 1356 |
+
//
|
| 1357 |
+
// Given an operator schema without an overload name: aten::op(...
|
| 1358 |
+
//
|
| 1359 |
+
// Use ATEN_FN(op) to get an unambiguous *function* version of the operator.
|
| 1360 |
+
//
|
| 1361 |
+
// There is some interesting behavior for out= operations.
|
| 1362 |
+
// ATEN_FN2(sin, out) gives a function that is *faithful* to the schema;
|
| 1363 |
+
// that is, the order of arguments is exactly what it looks like in the schema.
|
| 1364 |
+
|
| 1365 |
+
#define ATEN_FN2(op_name, overload) at::_ops::op_name##_##overload::call
|
| 1366 |
+
#define ATEN_FN(op_name) at::_ops::op_name::call
|
| 1367 |
+
|
| 1368 |
+
// Separately, ATEN_OP(op) and ATEN_OP2(op, overload) define a class containing compile-time
|
| 1369 |
+
// metadata about a given aten operator.
|
| 1370 |
+
// Notable data on the class includes:
|
| 1371 |
+
// - ATEN_OP2(add, Tensor)::name // returns the string name: "add"
|
| 1372 |
+
// - ATEN_OP2(add, Tensor)::overload_name // returns the string overload name: "Tensor"
|
| 1373 |
+
// - ATEN_OP2(add, Tensor)::schema // returns the C++ schema type: at::Tensor (const at::Tensor &, const at::Tensor &, const at::Scalar &)
|
| 1374 |
+
// - ATEN_OP2(add, Tensor)::schema_str // returns the string jit type: "add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor"
|
| 1375 |
+
|
| 1376 |
+
#define ATEN_OP2(op_name, overload) at::_ops::op_name##_##overload
|
| 1377 |
+
#define ATEN_OP(op_name) at::_ops::op_name
|
| 1378 |
+
|
| 1379 |
+
// WARNING: Please do not call any of the ops in the _ops namespace directly.
|
| 1380 |
+
// Use the ATEN_FN macros. We do not guarantee stability of the naming
|
| 1381 |
+
// scheme for the functions in at::_ops
|
| 1382 |
+
|
| 1383 |
+
// See Note [The ATen Operators API] for details of the at::_ops namespace
|
| 1384 |
+
|
| 1385 |
+
namespace at {
|
| 1386 |
+
namespace _ops {
|
| 1387 |
+
|
| 1388 |
+
} // namespace _ops
|
| 1389 |
+
} // namespace at
|
lib/python3.12/site-packages/torch/include/ATen/PTThreadPool.h
ADDED
|
@@ -0,0 +1,17 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <ATen/Parallel.h>
|
| 4 |
+
#include <c10/core/thread_pool.h>
|
| 5 |
+
|
| 6 |
+
namespace at {
|
| 7 |
+
|
| 8 |
+
class TORCH_API PTThreadPool : public c10::ThreadPool {
|
| 9 |
+
public:
|
| 10 |
+
explicit PTThreadPool(int pool_size, int numa_node_id = -1)
|
| 11 |
+
: c10::ThreadPool(pool_size, numa_node_id, []() {
|
| 12 |
+
c10::setThreadName("PTThreadPool");
|
| 13 |
+
at::init_num_threads();
|
| 14 |
+
}) {}
|
| 15 |
+
};
|
| 16 |
+
|
| 17 |
+
} // namespace at
|
lib/python3.12/site-packages/torch/include/ATen/PadNd.h
ADDED
|
@@ -0,0 +1,14 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
#include <c10/util/Exception.h>
|
| 3 |
+
#include <c10/util/string_view.h>
|
| 4 |
+
|
| 5 |
+
namespace at {
|
| 6 |
+
|
| 7 |
+
enum class padding_mode {
|
| 8 |
+
reflect,
|
| 9 |
+
replicate,
|
| 10 |
+
circular,
|
| 11 |
+
constant,
|
| 12 |
+
};
|
| 13 |
+
|
| 14 |
+
} // namespace at
|
lib/python3.12/site-packages/torch/include/ATen/PythonTorchFunctionTLS.h
ADDED
|
@@ -0,0 +1,36 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <c10/core/SafePyObject.h>
|
| 4 |
+
#include <c10/macros/Macros.h>
|
| 5 |
+
|
| 6 |
+
namespace at::impl {
|
| 7 |
+
|
| 8 |
+
enum TorchFunctionDisabledState { ENABLED, SUBCLASSES_DISABLED, ALL_DISABLED };
|
| 9 |
+
|
| 10 |
+
struct TORCH_API PythonTorchFunctionTLS {
|
| 11 |
+
static void set_disabled_state(TorchFunctionDisabledState disabled_state_);
|
| 12 |
+
static TorchFunctionDisabledState get_disabled_state();
|
| 13 |
+
|
| 14 |
+
static void push_onto_stack(std::shared_ptr<SafePyObject> mode);
|
| 15 |
+
static const std::shared_ptr<SafePyObject> pop_stack();
|
| 16 |
+
static const std::shared_ptr<SafePyObject>& get_stack_at(int64_t idx);
|
| 17 |
+
static int64_t stack_len();
|
| 18 |
+
|
| 19 |
+
static const PythonTorchFunctionTLS& get_state();
|
| 20 |
+
static void set_state(const PythonTorchFunctionTLS& state);
|
| 21 |
+
|
| 22 |
+
private:
|
| 23 |
+
// The mode TLS is split into
|
| 24 |
+
// - disabled_state, which says which part of torch function are disabled
|
| 25 |
+
// - stack_, which is a vector of modes representing the stack of user
|
| 26 |
+
// defined modes
|
| 27 |
+
TorchFunctionDisabledState disabled_state_ =
|
| 28 |
+
TorchFunctionDisabledState::ENABLED;
|
| 29 |
+
std::vector<std::shared_ptr<c10::SafePyObject>> stack_;
|
| 30 |
+
};
|
| 31 |
+
|
| 32 |
+
TORCH_API bool torch_function_mode_enabled();
|
| 33 |
+
|
| 34 |
+
TORCH_API bool torch_function_all_disabled();
|
| 35 |
+
|
| 36 |
+
} // namespace at::impl
|
lib/python3.12/site-packages/torch/include/ATen/ScalarOps.h
ADDED
|
@@ -0,0 +1,53 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <ATen/Tensor.h>
|
| 4 |
+
#include <c10/core/Scalar.h>
|
| 5 |
+
|
| 6 |
+
#ifndef AT_PER_OPERATOR_HEADERS
|
| 7 |
+
#include <ATen/Functions.h>
|
| 8 |
+
#else
|
| 9 |
+
#include <ATen/ops/scalar_tensor.h>
|
| 10 |
+
#endif
|
| 11 |
+
|
| 12 |
+
namespace at::detail {
|
| 13 |
+
// When filling a number to 1-element CPU tensor, we want to skip
|
| 14 |
+
// everything but manipulate data ptr directly.
|
| 15 |
+
// Ideally this fast pass should be implemented in TensorIterator,
|
| 16 |
+
// but we also want to skip compute_types which in not avoidable
|
| 17 |
+
// in TensorIterator for now.
|
| 18 |
+
Tensor& scalar_fill(Tensor& self, const Scalar& value);
|
| 19 |
+
TORCH_API Tensor scalar_tensor_static(
|
| 20 |
+
const Scalar& s,
|
| 21 |
+
std::optional<ScalarType> dtype_opt,
|
| 22 |
+
std::optional<Device> device_opt);
|
| 23 |
+
} // namespace at::detail
|
| 24 |
+
|
| 25 |
+
// This is in the c10 namespace because we use ADL to find the functions in it.
|
| 26 |
+
namespace c10 {
|
| 27 |
+
|
| 28 |
+
// FIXME: this should be (and was) Scalar::toTensor, but there is currently no
|
| 29 |
+
// way to implement this without going through Derived Types (which are not part
|
| 30 |
+
// of core).
|
| 31 |
+
inline at::Tensor scalar_to_tensor(
|
| 32 |
+
const Scalar& s,
|
| 33 |
+
const Device device = at::kCPU) {
|
| 34 |
+
// This is the fast track we have for CPU scalar tensors.
|
| 35 |
+
if (device == at::kCPU) {
|
| 36 |
+
return at::detail::scalar_tensor_static(s, s.type(), at::kCPU);
|
| 37 |
+
}
|
| 38 |
+
return at::scalar_tensor(s, at::device(device).dtype(s.type()));
|
| 39 |
+
}
|
| 40 |
+
|
| 41 |
+
} // namespace c10
|
| 42 |
+
|
| 43 |
+
namespace at::native {
|
| 44 |
+
|
| 45 |
+
inline Tensor wrapped_scalar_tensor(
|
| 46 |
+
const Scalar& scalar,
|
| 47 |
+
const Device device = at::kCPU) {
|
| 48 |
+
auto tensor = scalar_to_tensor(scalar, device);
|
| 49 |
+
tensor.unsafeGetTensorImpl()->set_wrapped_number(true);
|
| 50 |
+
return tensor;
|
| 51 |
+
}
|
| 52 |
+
|
| 53 |
+
} // namespace at::native
|
lib/python3.12/site-packages/torch/include/ATen/ScalarType.h
ADDED
|
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
#include <ATen/core/ATenGeneral.h> // for BC reasons
|
| 3 |
+
#include <c10/core/Backend.h>
|
| 4 |
+
#include <c10/core/ScalarType.h>
|
lib/python3.12/site-packages/torch/include/ATen/SmallVector.h
ADDED
|
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
#include <c10/util/SmallVector.h>
|
lib/python3.12/site-packages/torch/include/ATen/SparseCsrTensorImpl.h
ADDED
|
@@ -0,0 +1,206 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <ATen/Tensor.h>
|
| 4 |
+
#include <c10/core/TensorImpl.h>
|
| 5 |
+
#include <c10/core/impl/TorchDispatchModeTLS.h>
|
| 6 |
+
#include <c10/util/Exception.h>
|
| 7 |
+
namespace at {
|
| 8 |
+
|
| 9 |
+
// Struct implementing a sparse CSR tensor. It uses three 1-D tensors for
|
| 10 |
+
// denoting the data: `crow_indices_`, `col_indices_` and `values_`.
|
| 11 |
+
// The `crow_indices_` tensor is a integer tensor of shape `(size(0) + 1)`
|
| 12 |
+
// that represents the compressed row indices of the CSR tensor. The
|
| 13 |
+
// `col_indices_` tensor is an integer tensor of shape `(nnz())`
|
| 14 |
+
// that explicitly stores the column indices of each value of the sparse
|
| 15 |
+
// tensor. The `values_` tensor can be of any pytorch-supported data type
|
| 16 |
+
// and has shape `(nnz())`.
|
| 17 |
+
//
|
| 18 |
+
// Since the main advantage of the CSR format over the COO format is speed of
|
| 19 |
+
// computation, care must be taken to facilitate smooth interfacing of
|
| 20 |
+
// these data structures with optimized libraries such as MKL and MAGMA.
|
| 21 |
+
// Since the MKL interface for pytorch currently uses indexing with int32
|
| 22 |
+
// type, it is important to make sure that the `crow_indices` and `col_indices`
|
| 23 |
+
// are of type int32 when calling MKL routines such as SPMM or SPMV.
|
| 24 |
+
//
|
| 25 |
+
// If not calling MKL, it should be alright to use 64 bit integer tensors
|
| 26 |
+
// for indexing.
|
| 27 |
+
struct TORCH_API SparseCsrTensorImpl : public TensorImpl {
|
| 28 |
+
Tensor crow_indices_;
|
| 29 |
+
Tensor col_indices_;
|
| 30 |
+
Tensor values_;
|
| 31 |
+
Layout layout_;
|
| 32 |
+
|
| 33 |
+
public:
|
| 34 |
+
explicit SparseCsrTensorImpl(
|
| 35 |
+
at::DispatchKeySet,
|
| 36 |
+
at::Device device,
|
| 37 |
+
Layout layout,
|
| 38 |
+
const caffe2::TypeMeta);
|
| 39 |
+
|
| 40 |
+
void resize_(int64_t nnz, IntArrayRef size);
|
| 41 |
+
void resize_and_clear_(
|
| 42 |
+
int64_t sparse_dim,
|
| 43 |
+
int64_t dense_dim,
|
| 44 |
+
IntArrayRef size);
|
| 45 |
+
void resize_as_sparse_compressed_tensor_(const Tensor& src);
|
| 46 |
+
void set_member_tensors(
|
| 47 |
+
const Tensor& crow_indices,
|
| 48 |
+
const Tensor& col_indices,
|
| 49 |
+
const Tensor& values,
|
| 50 |
+
c10::SymIntArrayRef size);
|
| 51 |
+
void set_member_tensors(
|
| 52 |
+
const Tensor& crow_indices,
|
| 53 |
+
const Tensor& col_indices,
|
| 54 |
+
const Tensor& values,
|
| 55 |
+
IntArrayRef size);
|
| 56 |
+
const Tensor& compressed_indices() const {
|
| 57 |
+
return crow_indices_;
|
| 58 |
+
}
|
| 59 |
+
const Tensor& plain_indices() const {
|
| 60 |
+
return col_indices_;
|
| 61 |
+
}
|
| 62 |
+
const Tensor& values() const {
|
| 63 |
+
return values_;
|
| 64 |
+
}
|
| 65 |
+
int64_t nnz() {
|
| 66 |
+
return col_indices_.size(-1);
|
| 67 |
+
}
|
| 68 |
+
|
| 69 |
+
inline int64_t batch_dim() const noexcept {
|
| 70 |
+
return crow_indices_.dim() - 1;
|
| 71 |
+
}
|
| 72 |
+
|
| 73 |
+
inline int64_t sparse_dim() const noexcept {
|
| 74 |
+
return 2;
|
| 75 |
+
}
|
| 76 |
+
|
| 77 |
+
inline int64_t dense_dim() const noexcept {
|
| 78 |
+
return values_.dim() - batch_dim() - block_dim() - 1;
|
| 79 |
+
}
|
| 80 |
+
|
| 81 |
+
private:
|
| 82 |
+
inline int64_t block_dim() const noexcept {
|
| 83 |
+
return (layout_ == kSparseBsr || layout_ == kSparseBsc ? 2 : 0);
|
| 84 |
+
}
|
| 85 |
+
|
| 86 |
+
protected:
|
| 87 |
+
IntArrayRef strides_custom() const override;
|
| 88 |
+
SymIntArrayRef sym_strides_custom() const override;
|
| 89 |
+
bool is_contiguous_custom(MemoryFormat) const override;
|
| 90 |
+
|
| 91 |
+
public:
|
| 92 |
+
void set_size(int64_t dim, int64_t new_size) override;
|
| 93 |
+
void set_stride(int64_t dim, int64_t new_stride) override;
|
| 94 |
+
void set_storage_offset(int64_t storage_offset) override;
|
| 95 |
+
Layout layout_impl() const override {
|
| 96 |
+
return layout_;
|
| 97 |
+
}
|
| 98 |
+
void set_layout(Layout layout) {
|
| 99 |
+
switch (layout) {
|
| 100 |
+
case kSparseCsr:
|
| 101 |
+
case kSparseCsc:
|
| 102 |
+
case kSparseBsr:
|
| 103 |
+
case kSparseBsc:
|
| 104 |
+
layout_ = layout;
|
| 105 |
+
break;
|
| 106 |
+
default:
|
| 107 |
+
TORCH_CHECK(false, "unsupported layout ", layout);
|
| 108 |
+
}
|
| 109 |
+
}
|
| 110 |
+
|
| 111 |
+
template <typename VariableVersion>
|
| 112 |
+
c10::intrusive_ptr<TensorImpl> shallow_copy_and_detach_core(
|
| 113 |
+
VariableVersion&& version_counter,
|
| 114 |
+
bool allow_tensor_metadata_change) const {
|
| 115 |
+
const auto mode_stack_len = c10::impl::TorchDispatchModeTLS::stack_len();
|
| 116 |
+
c10::impl::PyInterpreter&& interpreter = nullptr;
|
| 117 |
+
if (mode_stack_len > 0 &&
|
| 118 |
+
!c10::impl::tls_is_dispatch_key_excluded(DispatchKey::Python)) {
|
| 119 |
+
const auto& cur_torch_dispatch_mode_state =
|
| 120 |
+
c10::impl::TorchDispatchModeTLS::get_stack_at(mode_stack_len - 1);
|
| 121 |
+
interpreter = cur_torch_dispatch_mode_state->pyinterpreter();
|
| 122 |
+
} else if (
|
| 123 |
+
key_set_.has(DispatchKey::Python) &&
|
| 124 |
+
!c10::impl::tls_is_dispatch_key_excluded(DispatchKey::Python)) {
|
| 125 |
+
interpreter = pyobj_slot_.load_pyobj_interpreter();
|
| 126 |
+
} else {
|
| 127 |
+
// otherwise just copy the SparseTensorImpl and not the PyObject.
|
| 128 |
+
auto impl = c10::make_intrusive<SparseCsrTensorImpl>(
|
| 129 |
+
key_set(), device(), layout_impl(), dtype());
|
| 130 |
+
copy_tensor_metadata(
|
| 131 |
+
/*src_sparse_impl=*/this,
|
| 132 |
+
/*dest_sparse_impl=*/impl.get(),
|
| 133 |
+
/*version_counter=*/version_counter,
|
| 134 |
+
/*allow_tensor_metadata_change=*/allow_tensor_metadata_change);
|
| 135 |
+
impl->refresh_numel();
|
| 136 |
+
return impl;
|
| 137 |
+
}
|
| 138 |
+
auto r = interpreter->detach(this);
|
| 139 |
+
r->set_version_counter(std::forward<VariableVersion>(version_counter));
|
| 140 |
+
r->set_allow_tensor_metadata_change(allow_tensor_metadata_change);
|
| 141 |
+
return r;
|
| 142 |
+
}
|
| 143 |
+
|
| 144 |
+
/**
|
| 145 |
+
* Return a TensorImpl that is a shallow-copy of this TensorImpl.
|
| 146 |
+
*
|
| 147 |
+
* For usage of `version_counter` and `allow_tensor_metadata_change`,
|
| 148 |
+
* see NOTE [ TensorImpl Shallow-Copying ].
|
| 149 |
+
*/
|
| 150 |
+
c10::intrusive_ptr<TensorImpl> shallow_copy_and_detach(
|
| 151 |
+
const c10::VariableVersion& version_counter,
|
| 152 |
+
bool allow_tensor_metadata_change) const override {
|
| 153 |
+
return shallow_copy_and_detach_core(
|
| 154 |
+
version_counter, allow_tensor_metadata_change);
|
| 155 |
+
}
|
| 156 |
+
|
| 157 |
+
/**
|
| 158 |
+
* Return a TensorImpl that is a shallow-copy of this TensorImpl.
|
| 159 |
+
*
|
| 160 |
+
* For usage of `version_counter` and `allow_tensor_metadata_change`,
|
| 161 |
+
* see NOTE [ TensorImpl Shallow-Copying ].
|
| 162 |
+
*/
|
| 163 |
+
c10::intrusive_ptr<TensorImpl> shallow_copy_and_detach(
|
| 164 |
+
c10::VariableVersion&& version_counter,
|
| 165 |
+
bool allow_tensor_metadata_change) const override {
|
| 166 |
+
return shallow_copy_and_detach_core(
|
| 167 |
+
std::move(version_counter), allow_tensor_metadata_change);
|
| 168 |
+
}
|
| 169 |
+
|
| 170 |
+
private:
|
| 171 |
+
explicit SparseCsrTensorImpl(
|
| 172 |
+
at::DispatchKeySet key_set,
|
| 173 |
+
const caffe2::TypeMeta data_type,
|
| 174 |
+
at::Tensor crow_indices,
|
| 175 |
+
at::Tensor col_indices,
|
| 176 |
+
at::Tensor values,
|
| 177 |
+
at::Layout layout);
|
| 178 |
+
|
| 179 |
+
const char* tensorimpl_type_name() const override;
|
| 180 |
+
|
| 181 |
+
/**
|
| 182 |
+
* Copy the tensor metadata fields (e.g. sizes / strides / storage pointer /
|
| 183 |
+
* storage_offset) from one TensorImpl to another TensorImpl.
|
| 184 |
+
*
|
| 185 |
+
* For usage of `version_counter` and `allow_tensor_metadata_change`, see NOTE
|
| 186 |
+
* [ TensorImpl Shallow-Copying ].
|
| 187 |
+
*/
|
| 188 |
+
static void copy_tensor_metadata(
|
| 189 |
+
const SparseCsrTensorImpl* src_sparse_impl,
|
| 190 |
+
SparseCsrTensorImpl* dest_sparse_impl,
|
| 191 |
+
c10::VariableVersion version_counter,
|
| 192 |
+
bool allow_tensor_metadata_change) {
|
| 193 |
+
TensorImpl::copy_tensor_metadata(
|
| 194 |
+
src_sparse_impl,
|
| 195 |
+
dest_sparse_impl,
|
| 196 |
+
std::move(version_counter),
|
| 197 |
+
allow_tensor_metadata_change);
|
| 198 |
+
|
| 199 |
+
// Sparse-specific fields
|
| 200 |
+
dest_sparse_impl->crow_indices_ = src_sparse_impl->compressed_indices();
|
| 201 |
+
dest_sparse_impl->col_indices_ = src_sparse_impl->plain_indices();
|
| 202 |
+
dest_sparse_impl->values_ = src_sparse_impl->values();
|
| 203 |
+
dest_sparse_impl->layout_ = src_sparse_impl->layout_impl();
|
| 204 |
+
}
|
| 205 |
+
};
|
| 206 |
+
} // namespace at
|
lib/python3.12/site-packages/torch/include/ATen/SparseTensorImpl.h
ADDED
|
@@ -0,0 +1,421 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <ATen/Tensor.h>
|
| 4 |
+
#include <c10/core/TensorImpl.h>
|
| 5 |
+
#include <c10/core/impl/TorchDispatchModeTLS.h>
|
| 6 |
+
#include <c10/util/Exception.h>
|
| 7 |
+
#include <c10/util/irange.h>
|
| 8 |
+
|
| 9 |
+
#ifndef AT_PER_OPERATOR_HEADERS
|
| 10 |
+
#include <ATen/Functions.h>
|
| 11 |
+
#else
|
| 12 |
+
#include <ATen/ops/empty.h>
|
| 13 |
+
#include <ATen/ops/resize.h>
|
| 14 |
+
#endif
|
| 15 |
+
|
| 16 |
+
namespace at {
|
| 17 |
+
struct TORCH_API SparseTensorImpl : public TensorImpl {
|
| 18 |
+
// Stored in COO format, indices + values.
|
| 19 |
+
|
| 20 |
+
// INVARIANTS:
|
| 21 |
+
// sparse_dim: range [0, len(shape)]; sparse_dim + dense_dim = len(shape)
|
| 22 |
+
// dense_dim : range [0, len(shape)]; sparse_dim + dense_dim = len(shape)
|
| 23 |
+
// _indices.shape: dimensionality: 2, shape: (sparse_dim, nnz)
|
| 24 |
+
// _values.shape: dimensionality: 1 + dense_dim. shape: (nnz,
|
| 25 |
+
// shape[sparse_dim:])
|
| 26 |
+
|
| 27 |
+
int64_t sparse_dim_ = 0; // number of sparse dimensions
|
| 28 |
+
int64_t dense_dim_ = 0; // number of dense dimensions
|
| 29 |
+
|
| 30 |
+
Tensor indices_; // always a LongTensor
|
| 31 |
+
Tensor values_;
|
| 32 |
+
|
| 33 |
+
// A sparse tensor is 'coalesced' if every index occurs at most once in
|
| 34 |
+
// the indices tensor, and the indices are in sorted order. (This means
|
| 35 |
+
// that it is very easy to convert a coalesced tensor to CSR format: you
|
| 36 |
+
// need only compute CSR format indices.)
|
| 37 |
+
//
|
| 38 |
+
// Most math operations can only be performed on coalesced sparse tensors,
|
| 39 |
+
// because many algorithms proceed by merging two sorted lists (of indices).
|
| 40 |
+
bool coalesced_ = false;
|
| 41 |
+
|
| 42 |
+
// compute_numel with integer multiplication overflow check, see gh-57542
|
| 43 |
+
void refresh_numel() {
|
| 44 |
+
TensorImpl::safe_refresh_numel();
|
| 45 |
+
}
|
| 46 |
+
|
| 47 |
+
public:
|
| 48 |
+
// Public for now...
|
| 49 |
+
explicit SparseTensorImpl(at::DispatchKeySet, const caffe2::TypeMeta);
|
| 50 |
+
|
| 51 |
+
void release_resources() override;
|
| 52 |
+
|
| 53 |
+
int64_t nnz() const {
|
| 54 |
+
return values_.size(0);
|
| 55 |
+
}
|
| 56 |
+
|
| 57 |
+
c10::SymInt sym_nnz() const {
|
| 58 |
+
return values_.sym_size(0);
|
| 59 |
+
}
|
| 60 |
+
int64_t sparse_dim() const {
|
| 61 |
+
return sparse_dim_;
|
| 62 |
+
}
|
| 63 |
+
int64_t dense_dim() const {
|
| 64 |
+
return dense_dim_;
|
| 65 |
+
}
|
| 66 |
+
bool coalesced() const {
|
| 67 |
+
return coalesced_;
|
| 68 |
+
}
|
| 69 |
+
Tensor indices() const {
|
| 70 |
+
return indices_;
|
| 71 |
+
}
|
| 72 |
+
Tensor values() const {
|
| 73 |
+
return values_;
|
| 74 |
+
}
|
| 75 |
+
|
| 76 |
+
void set_size(int64_t dim, int64_t new_size) override;
|
| 77 |
+
void set_stride(int64_t dim, int64_t new_stride) override;
|
| 78 |
+
void set_storage_offset(int64_t storage_offset) override;
|
| 79 |
+
|
| 80 |
+
#ifdef DEBUG
|
| 81 |
+
bool has_storage() const override;
|
| 82 |
+
#endif
|
| 83 |
+
|
| 84 |
+
// WARNING: This function does NOT preserve invariants of sparse_dim/dense_dim
|
| 85 |
+
// with respect to indices and values
|
| 86 |
+
void raw_resize_(int64_t sparse_dim, int64_t dense_dim, IntArrayRef size) {
|
| 87 |
+
TORCH_CHECK(
|
| 88 |
+
allow_tensor_metadata_change(),
|
| 89 |
+
"raw_resize_ ",
|
| 90 |
+
err_msg_tensor_metadata_change_not_allowed);
|
| 91 |
+
TORCH_CHECK(
|
| 92 |
+
!has_symbolic_sizes_strides_,
|
| 93 |
+
"raw_resize_ called on tensor with symbolic shape")
|
| 94 |
+
set_sizes_and_strides(size, std::vector<int64_t>(size.size()));
|
| 95 |
+
sparse_dim_ = sparse_dim;
|
| 96 |
+
dense_dim_ = dense_dim;
|
| 97 |
+
refresh_numel();
|
| 98 |
+
}
|
| 99 |
+
|
| 100 |
+
// NOTE: This function preserves invariants of sparse_dim/dense_dim with
|
| 101 |
+
// respect to indices and values.
|
| 102 |
+
//
|
| 103 |
+
// NOTE: This function supports the following cases:
|
| 104 |
+
// 1. When we keep the number of dense dimensions unchanged, and NOT shrinking
|
| 105 |
+
// the size of any of the dense dimensions.
|
| 106 |
+
// 2. When we keep the number of sparse dimensions unchanged, and NOT
|
| 107 |
+
// shrinking the size of any of the sparse dimensions.
|
| 108 |
+
// 3. When the sparse tensor has zero nnz, in which case we are free to change
|
| 109 |
+
// the shapes of both its sparse and dense dimensions.
|
| 110 |
+
//
|
| 111 |
+
// This function DOESN'T support (and will throw an error) the following
|
| 112 |
+
// cases:
|
| 113 |
+
// 1. When we attempt to change the number of sparse dimensions on a non-empty
|
| 114 |
+
// sparse tensor (such an operation will invalidate the indices stored).
|
| 115 |
+
// 2. When we attempt to change the number of dense dimensions on a non-empty
|
| 116 |
+
// sparse tensor (such an operation will behave differently from an equivalent
|
| 117 |
+
// dense tensor's resize method, and for API consistency we don't support it).
|
| 118 |
+
// 3. When we attempt to shrink the size of any of the dense dimensions on a
|
| 119 |
+
// non-empty sparse tensor (such an operation will behave differently from an
|
| 120 |
+
// equivalent dense tensor's resize method, and for API consistency we don't
|
| 121 |
+
// support it).
|
| 122 |
+
// 4. When we attempt to shrink the size of any of the sparse dimensions on a
|
| 123 |
+
// non-empty sparse tensor (this could make some of the stored indices
|
| 124 |
+
// out-of-bound and thus unsafe).
|
| 125 |
+
template <typename T>
|
| 126 |
+
void _resize_(int64_t sparse_dim, int64_t dense_dim, ArrayRef<T> size) {
|
| 127 |
+
TORCH_CHECK(
|
| 128 |
+
allow_tensor_metadata_change(),
|
| 129 |
+
"resize_ ",
|
| 130 |
+
err_msg_tensor_metadata_change_not_allowed);
|
| 131 |
+
TORCH_CHECK(
|
| 132 |
+
!has_symbolic_sizes_strides_,
|
| 133 |
+
"resize_ called on tensor with symbolic shape")
|
| 134 |
+
TORCH_CHECK(
|
| 135 |
+
sparse_dim + dense_dim == static_cast<int64_t>(size.size()),
|
| 136 |
+
"number of dimensions must be sparse_dim (",
|
| 137 |
+
sparse_dim,
|
| 138 |
+
") + dense_dim (",
|
| 139 |
+
dense_dim,
|
| 140 |
+
"), but got ",
|
| 141 |
+
size.size());
|
| 142 |
+
if (nnz() > 0) {
|
| 143 |
+
[[maybe_unused]] auto constexpr alt_options_msg =
|
| 144 |
+
"You could try the following options:\n\
|
| 145 |
+
1. If you need an empty sparse tensor of this size, call `x = torch.sparse_coo_tensor(size)`.\n\
|
| 146 |
+
2. If you need to resize this tensor, you have the following options:\n\
|
| 147 |
+
1. For both sparse and dense dimensions, keep the number of them constant and the size of them non-shrinking, and then try the same call again.\n\
|
| 148 |
+
2. Or, create a new sparse tensor with the correct indices and values from this sparse tensor.";
|
| 149 |
+
|
| 150 |
+
TORCH_CHECK(
|
| 151 |
+
sparse_dim == sparse_dim_,
|
| 152 |
+
"changing the number of sparse dimensions (from ",
|
| 153 |
+
sparse_dim_,
|
| 154 |
+
" to ",
|
| 155 |
+
sparse_dim,
|
| 156 |
+
") on a non-empty sparse tensor is not supported.\n",
|
| 157 |
+
alt_options_msg);
|
| 158 |
+
|
| 159 |
+
TORCH_CHECK(
|
| 160 |
+
dense_dim == dense_dim_,
|
| 161 |
+
"changing the number of dense dimensions (from ",
|
| 162 |
+
dense_dim_,
|
| 163 |
+
" to ",
|
| 164 |
+
dense_dim,
|
| 165 |
+
") on a non-empty sparse tensor is not supported.\n",
|
| 166 |
+
alt_options_msg);
|
| 167 |
+
|
| 168 |
+
bool shrinking_sparse_dims = false;
|
| 169 |
+
bool shrinking_dense_dim = false;
|
| 170 |
+
auto sparse_size_original = generic_sizes<T>().slice(0, sparse_dim);
|
| 171 |
+
auto sparse_size_new = size.slice(0, sparse_dim);
|
| 172 |
+
for (const auto i : c10::irange(sparse_dim)) {
|
| 173 |
+
if (sparse_size_new[i] < sparse_size_original[i]) {
|
| 174 |
+
shrinking_sparse_dims = true;
|
| 175 |
+
break;
|
| 176 |
+
}
|
| 177 |
+
}
|
| 178 |
+
auto dense_size_original = generic_sizes<T>().slice(sparse_dim);
|
| 179 |
+
auto dense_size_new = size.slice(sparse_dim);
|
| 180 |
+
for (const auto i : c10::irange(dense_dim)) {
|
| 181 |
+
if (dense_size_new[i] < dense_size_original[i]) {
|
| 182 |
+
shrinking_dense_dim = true;
|
| 183 |
+
break;
|
| 184 |
+
}
|
| 185 |
+
}
|
| 186 |
+
|
| 187 |
+
TORCH_CHECK(
|
| 188 |
+
!shrinking_sparse_dims,
|
| 189 |
+
"shrinking the size of sparse dimensions (from ",
|
| 190 |
+
sparse_size_original,
|
| 191 |
+
" to ",
|
| 192 |
+
sparse_size_new,
|
| 193 |
+
") on a non-empty sparse tensor is not supported.\n",
|
| 194 |
+
alt_options_msg);
|
| 195 |
+
|
| 196 |
+
TORCH_CHECK(
|
| 197 |
+
!shrinking_dense_dim,
|
| 198 |
+
"shrinking the size of dense dimensions (from ",
|
| 199 |
+
dense_size_original,
|
| 200 |
+
" to ",
|
| 201 |
+
dense_size_new,
|
| 202 |
+
") on a non-empty sparse tensor is not supported.\n",
|
| 203 |
+
alt_options_msg);
|
| 204 |
+
}
|
| 205 |
+
|
| 206 |
+
auto sizes_and_strides = generic_sizes<T>();
|
| 207 |
+
const bool size_equals_sizes = std::equal(
|
| 208 |
+
size.begin(),
|
| 209 |
+
size.end(),
|
| 210 |
+
sizes_and_strides.begin(),
|
| 211 |
+
sizes_and_strides.end());
|
| 212 |
+
if ((!size_equals_sizes) || (sparse_dim != sparse_dim_) ||
|
| 213 |
+
(dense_dim != dense_dim_)) {
|
| 214 |
+
auto nnz = at::symint::sizes<T>(values())[0];
|
| 215 |
+
std::vector<T> values_size = {nnz};
|
| 216 |
+
auto dense_size = size.slice(sparse_dim);
|
| 217 |
+
values_size.insert(
|
| 218 |
+
values_size.end(), dense_size.begin(), dense_size.end());
|
| 219 |
+
at::symint::resize_<T>(values_, values_size);
|
| 220 |
+
at::symint::resize_<T>(indices_, {T(sparse_dim), nnz});
|
| 221 |
+
}
|
| 222 |
+
|
| 223 |
+
if (!size_equals_sizes) {
|
| 224 |
+
set_sizes_and_strides(size, std::vector<T>(size.size()));
|
| 225 |
+
}
|
| 226 |
+
sparse_dim_ = sparse_dim;
|
| 227 |
+
dense_dim_ = dense_dim;
|
| 228 |
+
refresh_numel();
|
| 229 |
+
}
|
| 230 |
+
|
| 231 |
+
void resize_(int64_t sparse_dim, int64_t dense_dim, ArrayRef<int64_t> size) {
|
| 232 |
+
return _resize_(sparse_dim, dense_dim, size);
|
| 233 |
+
}
|
| 234 |
+
|
| 235 |
+
void resize_(
|
| 236 |
+
int64_t sparse_dim,
|
| 237 |
+
int64_t dense_dim,
|
| 238 |
+
ArrayRef<c10::SymInt> size) {
|
| 239 |
+
return _resize_(sparse_dim, dense_dim, size);
|
| 240 |
+
}
|
| 241 |
+
|
| 242 |
+
// NOTE: this function will resize the sparse tensor and also set `indices`
|
| 243 |
+
// and `values` to empty.
|
| 244 |
+
void resize_and_clear_(
|
| 245 |
+
int64_t sparse_dim,
|
| 246 |
+
int64_t dense_dim,
|
| 247 |
+
IntArrayRef size) {
|
| 248 |
+
TORCH_CHECK(
|
| 249 |
+
allow_tensor_metadata_change(),
|
| 250 |
+
"resize_and_clear_ ",
|
| 251 |
+
err_msg_tensor_metadata_change_not_allowed);
|
| 252 |
+
TORCH_CHECK(
|
| 253 |
+
!has_symbolic_sizes_strides_,
|
| 254 |
+
"resize_and_clear_ called on tensor with symbolic shape")
|
| 255 |
+
TORCH_CHECK(
|
| 256 |
+
sparse_dim + dense_dim == static_cast<int64_t>(size.size()),
|
| 257 |
+
"number of dimensions must be sparse_dim (",
|
| 258 |
+
sparse_dim,
|
| 259 |
+
") + dense_dim (",
|
| 260 |
+
dense_dim,
|
| 261 |
+
"), but got ",
|
| 262 |
+
size.size());
|
| 263 |
+
|
| 264 |
+
set_sizes_and_strides(size, std::vector<int64_t>(size.size()));
|
| 265 |
+
sparse_dim_ = sparse_dim;
|
| 266 |
+
dense_dim_ = dense_dim;
|
| 267 |
+
|
| 268 |
+
auto empty_indices = at::empty({sparse_dim, 0}, indices().options());
|
| 269 |
+
std::vector<int64_t> values_size = {0};
|
| 270 |
+
auto dense_size = sizes().slice(sparse_dim);
|
| 271 |
+
values_size.insert(values_size.end(), dense_size.begin(), dense_size.end());
|
| 272 |
+
auto empty_values = at::empty(values_size, values().options());
|
| 273 |
+
set_indices_and_values_unsafe(empty_indices, empty_values);
|
| 274 |
+
refresh_numel();
|
| 275 |
+
}
|
| 276 |
+
|
| 277 |
+
void set_coalesced(bool coalesced) {
|
| 278 |
+
TORCH_CHECK(
|
| 279 |
+
allow_tensor_metadata_change(),
|
| 280 |
+
"set_coalesced ",
|
| 281 |
+
err_msg_tensor_metadata_change_not_allowed);
|
| 282 |
+
coalesced_ = coalesced;
|
| 283 |
+
}
|
| 284 |
+
|
| 285 |
+
// NOTE: this function is only used internally and not exposed to Python
|
| 286 |
+
// frontend
|
| 287 |
+
void set_nnz_and_narrow(int64_t new_nnz) {
|
| 288 |
+
TORCH_CHECK(
|
| 289 |
+
allow_tensor_metadata_change(),
|
| 290 |
+
"set_nnz_and_narrow ",
|
| 291 |
+
err_msg_tensor_metadata_change_not_allowed);
|
| 292 |
+
AT_ASSERT(new_nnz <= nnz());
|
| 293 |
+
indices_ = indices_.narrow(1, 0, new_nnz);
|
| 294 |
+
values_ = values_.narrow(0, 0, new_nnz);
|
| 295 |
+
if (new_nnz < 2) {
|
| 296 |
+
coalesced_ = true;
|
| 297 |
+
}
|
| 298 |
+
}
|
| 299 |
+
|
| 300 |
+
// Takes indices and values and directly puts them into the sparse tensor, no
|
| 301 |
+
// copy. NOTE: this function is unsafe because it doesn't check whether any
|
| 302 |
+
// indices are out of boundaries of `sizes`, so it should ONLY be used where
|
| 303 |
+
// we know that the indices are guaranteed to be within bounds. This used to
|
| 304 |
+
// be called THSTensor_(_move) NB: This used to be able to avoid a refcount
|
| 305 |
+
// bump, but I was too lazy to make it happen
|
| 306 |
+
void set_indices_and_values_unsafe(
|
| 307 |
+
const Tensor& indices,
|
| 308 |
+
const Tensor& values);
|
| 309 |
+
|
| 310 |
+
template <typename VariableVersion>
|
| 311 |
+
c10::intrusive_ptr<TensorImpl> shallow_copy_and_detach_core(
|
| 312 |
+
VariableVersion&& version_counter,
|
| 313 |
+
bool allow_tensor_metadata_change) const {
|
| 314 |
+
const auto mode_stack_len = c10::impl::TorchDispatchModeTLS::stack_len();
|
| 315 |
+
c10::impl::PyInterpreter&& interpreter = nullptr;
|
| 316 |
+
if (mode_stack_len > 0 &&
|
| 317 |
+
!c10::impl::tls_is_dispatch_key_excluded(DispatchKey::Python)) {
|
| 318 |
+
const auto& cur_torch_dispatch_mode_state =
|
| 319 |
+
c10::impl::TorchDispatchModeTLS::get_stack_at(mode_stack_len - 1);
|
| 320 |
+
interpreter = cur_torch_dispatch_mode_state->pyinterpreter();
|
| 321 |
+
} else if (
|
| 322 |
+
key_set_.has(DispatchKey::Python) &&
|
| 323 |
+
!c10::impl::tls_is_dispatch_key_excluded(DispatchKey::Python)) {
|
| 324 |
+
interpreter = pyobj_slot_.load_pyobj_interpreter();
|
| 325 |
+
} else {
|
| 326 |
+
// otherwise just copy the SparseTensorImpl and not the PyObject.
|
| 327 |
+
auto impl = c10::make_intrusive<SparseTensorImpl>(key_set(), dtype());
|
| 328 |
+
copy_tensor_metadata(
|
| 329 |
+
/*src_sparse_impl=*/this,
|
| 330 |
+
/*dest_sparse_impl=*/impl.get(),
|
| 331 |
+
/*version_counter=*/version_counter,
|
| 332 |
+
/*allow_tensor_metadata_change=*/allow_tensor_metadata_change);
|
| 333 |
+
impl->refresh_numel();
|
| 334 |
+
return impl;
|
| 335 |
+
}
|
| 336 |
+
auto r = interpreter->detach(this);
|
| 337 |
+
r->set_version_counter(std::forward<VariableVersion>(version_counter));
|
| 338 |
+
r->set_allow_tensor_metadata_change(allow_tensor_metadata_change);
|
| 339 |
+
return r;
|
| 340 |
+
}
|
| 341 |
+
|
| 342 |
+
/**
|
| 343 |
+
* Return a TensorImpl that is a shallow-copy of this TensorImpl.
|
| 344 |
+
*
|
| 345 |
+
* For usage of `version_counter` and `allow_tensor_metadata_change`,
|
| 346 |
+
* see NOTE [ TensorImpl Shallow-Copying ].
|
| 347 |
+
*/
|
| 348 |
+
c10::intrusive_ptr<TensorImpl> shallow_copy_and_detach(
|
| 349 |
+
const c10::VariableVersion& version_counter,
|
| 350 |
+
bool allow_tensor_metadata_change) const override {
|
| 351 |
+
return shallow_copy_and_detach_core(
|
| 352 |
+
version_counter, allow_tensor_metadata_change);
|
| 353 |
+
}
|
| 354 |
+
|
| 355 |
+
/**
|
| 356 |
+
* Return a TensorImpl that is a shallow-copy of this TensorImpl.
|
| 357 |
+
*
|
| 358 |
+
* For usage of `version_counter` and `allow_tensor_metadata_change`,
|
| 359 |
+
* see NOTE [ TensorImpl Shallow-Copying ].
|
| 360 |
+
*/
|
| 361 |
+
c10::intrusive_ptr<TensorImpl> shallow_copy_and_detach(
|
| 362 |
+
c10::VariableVersion&& version_counter,
|
| 363 |
+
bool allow_tensor_metadata_change) const override {
|
| 364 |
+
return shallow_copy_and_detach_core(
|
| 365 |
+
std::move(version_counter), allow_tensor_metadata_change);
|
| 366 |
+
}
|
| 367 |
+
|
| 368 |
+
/**
|
| 369 |
+
* Shallow-copies data from another TensorImpl into this TensorImpl.
|
| 370 |
+
*
|
| 371 |
+
* For why this function doesn't check this TensorImpl's
|
| 372 |
+
* `allow_tensor_metadata_change_`, see NOTE [ TensorImpl Shallow-Copying ].
|
| 373 |
+
*/
|
| 374 |
+
void shallow_copy_from(const c10::intrusive_ptr<TensorImpl>& impl) override {
|
| 375 |
+
AT_ASSERT(has_compatible_shallow_copy_type(impl->key_set()));
|
| 376 |
+
auto sparse_impl = static_cast<const SparseTensorImpl*>(impl.get());
|
| 377 |
+
copy_tensor_metadata(
|
| 378 |
+
/*src_sparse_impl=*/sparse_impl,
|
| 379 |
+
/*dest_sparse_impl=*/this,
|
| 380 |
+
/*version_counter=*/version_counter(),
|
| 381 |
+
/*allow_tensor_metadata_change=*/allow_tensor_metadata_change());
|
| 382 |
+
refresh_numel();
|
| 383 |
+
}
|
| 384 |
+
|
| 385 |
+
private:
|
| 386 |
+
explicit SparseTensorImpl(
|
| 387 |
+
at::DispatchKeySet,
|
| 388 |
+
const caffe2::TypeMeta,
|
| 389 |
+
at::Tensor indices,
|
| 390 |
+
at::Tensor values);
|
| 391 |
+
|
| 392 |
+
/**
|
| 393 |
+
* Copy the tensor metadata fields (e.g. sizes / strides / storage pointer /
|
| 394 |
+
* storage_offset) from one TensorImpl to another TensorImpl.
|
| 395 |
+
*
|
| 396 |
+
* For usage of `version_counter` and `allow_tensor_metadata_change`, see NOTE
|
| 397 |
+
* [ TensorImpl Shallow-Copying ].
|
| 398 |
+
*/
|
| 399 |
+
static void copy_tensor_metadata(
|
| 400 |
+
const SparseTensorImpl* src_sparse_impl,
|
| 401 |
+
SparseTensorImpl* dest_sparse_impl,
|
| 402 |
+
c10::VariableVersion version_counter,
|
| 403 |
+
bool allow_tensor_metadata_change) {
|
| 404 |
+
TensorImpl::copy_tensor_metadata(
|
| 405 |
+
src_sparse_impl,
|
| 406 |
+
dest_sparse_impl,
|
| 407 |
+
std::move(version_counter),
|
| 408 |
+
allow_tensor_metadata_change);
|
| 409 |
+
|
| 410 |
+
// Sparse-specific fields
|
| 411 |
+
dest_sparse_impl->sparse_dim_ = src_sparse_impl->sparse_dim();
|
| 412 |
+
dest_sparse_impl->dense_dim_ = src_sparse_impl->dense_dim();
|
| 413 |
+
dest_sparse_impl->indices_ = src_sparse_impl->indices();
|
| 414 |
+
dest_sparse_impl->values_ = src_sparse_impl->values();
|
| 415 |
+
dest_sparse_impl->coalesced_ = src_sparse_impl->coalesced();
|
| 416 |
+
}
|
| 417 |
+
|
| 418 |
+
const char* tensorimpl_type_name() const override;
|
| 419 |
+
};
|
| 420 |
+
|
| 421 |
+
} // namespace at
|
lib/python3.12/site-packages/torch/include/ATen/Tensor.h
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <ATen/core/Tensor.h>
|
lib/python3.12/site-packages/torch/include/ATen/TensorGeometry.h
ADDED
|
@@ -0,0 +1,144 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <ATen/core/TensorBase.h>
|
| 4 |
+
#include <c10/core/WrapDimMinimal.h>
|
| 5 |
+
|
| 6 |
+
namespace at {
|
| 7 |
+
|
| 8 |
+
// Return if the tensor geometry represented by `sizes` and `strides` is
|
| 9 |
+
// contiguous Although we cache is_contiguous in tensor now, this is till useful
|
| 10 |
+
// because it allows checking if a particular geometry is contiguous without
|
| 11 |
+
// explicitly constructing a tensor, e.g., when you want to choose a kernel
|
| 12 |
+
// strategy based on whether a subgeometry is contiguous.
|
| 13 |
+
TORCH_API bool geometry_is_contiguous(IntArrayRef sizes, IntArrayRef strides);
|
| 14 |
+
|
| 15 |
+
struct TORCH_API TensorGeometry {
|
| 16 |
+
TensorGeometry() = default;
|
| 17 |
+
|
| 18 |
+
explicit TensorGeometry(c10::SymIntArrayRef sizes)
|
| 19 |
+
: sizes_(sizes.vec()),
|
| 20 |
+
strides_(sizes.size()),
|
| 21 |
+
has_symbolic_sizes_strides_(
|
| 22 |
+
!c10::asIntArrayRefSlowOpt(sizes).has_value()) {
|
| 23 |
+
int64_t dim = static_cast<int64_t>(sizes.size());
|
| 24 |
+
c10::SymInt expected_stride = 1;
|
| 25 |
+
for (int64_t i = dim - 1; i >= 0; i--) {
|
| 26 |
+
strides_[i] = expected_stride;
|
| 27 |
+
expected_stride *= sizes_[i];
|
| 28 |
+
}
|
| 29 |
+
numel_ = expected_stride;
|
| 30 |
+
}
|
| 31 |
+
|
| 32 |
+
explicit TensorGeometry(const TensorBase& t)
|
| 33 |
+
: sizes_(t.sym_sizes().vec()),
|
| 34 |
+
strides_(t.sym_strides().vec()),
|
| 35 |
+
storage_offset_(t.sym_storage_offset()),
|
| 36 |
+
numel_(t.sym_numel()),
|
| 37 |
+
has_symbolic_sizes_strides_(
|
| 38 |
+
t.unsafeGetTensorImpl()->has_symbolic_sizes_strides()) {}
|
| 39 |
+
|
| 40 |
+
// true if the tensor is contiguous
|
| 41 |
+
bool is_contiguous() const;
|
| 42 |
+
|
| 43 |
+
int64_t dim() const {
|
| 44 |
+
return static_cast<int64_t>(sizes_.size());
|
| 45 |
+
}
|
| 46 |
+
|
| 47 |
+
int64_t size(int64_t dim) const {
|
| 48 |
+
TORCH_INTERNAL_ASSERT(!has_symbolic_sizes_strides_);
|
| 49 |
+
dim = c10::maybe_wrap_dim(dim, this->dim());
|
| 50 |
+
return sizes_.at(static_cast<size_t>(dim)).as_int_unchecked();
|
| 51 |
+
}
|
| 52 |
+
c10::IntArrayRef sizes() const {
|
| 53 |
+
TORCH_INTERNAL_ASSERT(!has_symbolic_sizes_strides_);
|
| 54 |
+
return c10::asIntArrayRefUnchecked(sizes_);
|
| 55 |
+
}
|
| 56 |
+
int64_t stride(int64_t dim) const {
|
| 57 |
+
TORCH_INTERNAL_ASSERT(!has_symbolic_sizes_strides_);
|
| 58 |
+
dim = c10::maybe_wrap_dim(dim, this->dim());
|
| 59 |
+
return strides_.at(static_cast<size_t>(dim)).as_int_unchecked();
|
| 60 |
+
}
|
| 61 |
+
c10::IntArrayRef strides() const {
|
| 62 |
+
TORCH_INTERNAL_ASSERT(!has_symbolic_sizes_strides_);
|
| 63 |
+
return c10::asIntArrayRefUnchecked(strides_);
|
| 64 |
+
}
|
| 65 |
+
int64_t storage_offset() const {
|
| 66 |
+
TORCH_INTERNAL_ASSERT(!has_symbolic_sizes_strides_);
|
| 67 |
+
return storage_offset_.as_int_unchecked();
|
| 68 |
+
}
|
| 69 |
+
int64_t numel() const {
|
| 70 |
+
TORCH_INTERNAL_ASSERT(!has_symbolic_sizes_strides_);
|
| 71 |
+
return numel_.as_int_unchecked();
|
| 72 |
+
}
|
| 73 |
+
|
| 74 |
+
c10::SymInt sym_size(int64_t dim) const {
|
| 75 |
+
dim = c10::maybe_wrap_dim(dim, this->dim());
|
| 76 |
+
return sizes_.at(static_cast<size_t>(dim));
|
| 77 |
+
}
|
| 78 |
+
c10::SymIntArrayRef sym_sizes() const {
|
| 79 |
+
return sizes_;
|
| 80 |
+
}
|
| 81 |
+
c10::SymInt sym_stride(int64_t dim) const {
|
| 82 |
+
dim = c10::maybe_wrap_dim(dim, this->dim());
|
| 83 |
+
return strides_.at(static_cast<size_t>(dim));
|
| 84 |
+
}
|
| 85 |
+
c10::SymIntArrayRef sym_strides() const {
|
| 86 |
+
return strides_;
|
| 87 |
+
}
|
| 88 |
+
c10::SymInt sym_storage_offset() const {
|
| 89 |
+
return storage_offset_;
|
| 90 |
+
}
|
| 91 |
+
c10::SymInt sym_numel() const {
|
| 92 |
+
return numel_;
|
| 93 |
+
}
|
| 94 |
+
|
| 95 |
+
TensorGeometry transpose(int64_t dim0, int64_t dim1) {
|
| 96 |
+
TensorGeometry r = *this; // copy
|
| 97 |
+
TORCH_CHECK(
|
| 98 |
+
dim0 < dim(),
|
| 99 |
+
"transpose: dim0=",
|
| 100 |
+
dim0,
|
| 101 |
+
" out of range (dim=",
|
| 102 |
+
dim(),
|
| 103 |
+
")")
|
| 104 |
+
TORCH_CHECK(
|
| 105 |
+
dim1 < dim(),
|
| 106 |
+
"transpose: dim1=",
|
| 107 |
+
dim1,
|
| 108 |
+
" out of range (dim=",
|
| 109 |
+
dim(),
|
| 110 |
+
")")
|
| 111 |
+
std::swap(r.sizes_[dim0], r.sizes_[dim1]);
|
| 112 |
+
std::swap(r.strides_[dim0], r.strides_[dim1]);
|
| 113 |
+
return r;
|
| 114 |
+
}
|
| 115 |
+
|
| 116 |
+
std::vector<c10::SymInt>& mutable_sizes() {
|
| 117 |
+
return sizes_;
|
| 118 |
+
}
|
| 119 |
+
std::vector<c10::SymInt>& mutable_strides() {
|
| 120 |
+
return strides_;
|
| 121 |
+
}
|
| 122 |
+
c10::SymInt& mutable_storage_offset() {
|
| 123 |
+
return storage_offset_;
|
| 124 |
+
}
|
| 125 |
+
void recompute() {
|
| 126 |
+
// recalculate numel after a change
|
| 127 |
+
c10::SymInt numel = 1;
|
| 128 |
+
for (const auto& i : sizes_) {
|
| 129 |
+
numel = numel * i;
|
| 130 |
+
}
|
| 131 |
+
numel_ = std::move(numel);
|
| 132 |
+
has_symbolic_sizes_strides_ =
|
| 133 |
+
!c10::asIntArrayRefSlowOpt(sizes_).has_value();
|
| 134 |
+
}
|
| 135 |
+
|
| 136 |
+
private:
|
| 137 |
+
std::vector<c10::SymInt> sizes_;
|
| 138 |
+
std::vector<c10::SymInt> strides_;
|
| 139 |
+
c10::SymInt storage_offset_;
|
| 140 |
+
c10::SymInt numel_;
|
| 141 |
+
bool has_symbolic_sizes_strides_{false};
|
| 142 |
+
};
|
| 143 |
+
|
| 144 |
+
} // namespace at
|
lib/python3.12/site-packages/torch/include/ATen/TensorMeta.h
ADDED
|
@@ -0,0 +1,137 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <ATen/DimVector.h>
|
| 4 |
+
#include <ATen/core/Dimname.h>
|
| 5 |
+
#include <c10/core/TensorOptions.h>
|
| 6 |
+
#include <c10/util/strides.h>
|
| 7 |
+
|
| 8 |
+
namespace at {
|
| 9 |
+
|
| 10 |
+
class Tensor;
|
| 11 |
+
|
| 12 |
+
namespace impl {
|
| 13 |
+
|
| 14 |
+
// Use this to define the prototype for a meta function. There are two
|
| 15 |
+
// versions; one that takes one argument (just the operator name), or FUNC2
|
| 16 |
+
// variant that takes two arguments (operator name and overload name).
|
| 17 |
+
//
|
| 18 |
+
// Example usage:
|
| 19 |
+
//
|
| 20 |
+
// TORCH_META_FUNC2(add, Tensor) (
|
| 21 |
+
// const Tensor& self, const Tensor& other
|
| 22 |
+
// ) {
|
| 23 |
+
// ... compute sizes and options ...
|
| 24 |
+
// set_output(sizes, options);
|
| 25 |
+
// }
|
| 26 |
+
//
|
| 27 |
+
#define TORCH_META_FUNC(name) void structured_##name::meta
|
| 28 |
+
#define TORCH_META_FUNC2(name, overload) \
|
| 29 |
+
void structured_##name##_##overload::meta
|
| 30 |
+
|
| 31 |
+
// These are versions of TORCH_META_FUNC(2) that include a precompute_out struct
|
| 32 |
+
// as a return value. They should be used when the kernel in question has
|
| 33 |
+
// precomputed values declared in native_functions.yaml and the corresponding
|
| 34 |
+
// implementation should return an instance of the aforementioned struct.
|
| 35 |
+
#define TORCH_PRECOMPUTE_META_FUNC(name) \
|
| 36 |
+
structured_##name::meta_return_ty structured_##name::meta
|
| 37 |
+
#define TORCH_PRECOMPUTE_META_FUNC2(name, overload) \
|
| 38 |
+
structured_##name##_##overload::meta_return_ty \
|
| 39 |
+
structured_##name##_##overload::meta
|
| 40 |
+
|
| 41 |
+
// Use this to create a precompute struct in a meta function.
|
| 42 |
+
#define TORCH_PRECOMPUTE_STRUCT(name) structured_##name::precompute_out<>
|
| 43 |
+
#define TORCH_PRECOMPUTE_STRUCT2(name, overload) \
|
| 44 |
+
structured_##name##_##overload::precompute_out<>
|
| 45 |
+
|
| 46 |
+
// Use this to define the prototype for an implementation. This takes only
|
| 47 |
+
// one argument, which is the name of the dispatch key entry you're
|
| 48 |
+
// implementing.
|
| 49 |
+
//
|
| 50 |
+
// Example usage:
|
| 51 |
+
//
|
| 52 |
+
// TORCH_IMPL_FUNC(add_cpu) (
|
| 53 |
+
// Tensor& result, const Tensor& self, const Tensor& other
|
| 54 |
+
// ) {
|
| 55 |
+
// ... do the actual implementation ...
|
| 56 |
+
// }
|
| 57 |
+
//
|
| 58 |
+
#define TORCH_IMPL_FUNC(name) void structured_##name::impl
|
| 59 |
+
|
| 60 |
+
// Base class for all structured kernel classes. The set_output virtual
|
| 61 |
+
// method is varied depending whether or not the operator is
|
| 62 |
+
// functional/out/inplace, and could also be specialized for CPU/CUDA/etc
|
| 63 |
+
// (although presently it isn't).
|
| 64 |
+
//
|
| 65 |
+
// A notable subclass of this interface is TensorIteratorBase.
|
| 66 |
+
struct TORCH_API MetaBase {
|
| 67 |
+
MetaBase() = default;
|
| 68 |
+
MetaBase(const MetaBase&) = default;
|
| 69 |
+
MetaBase& operator=(const MetaBase&) = default;
|
| 70 |
+
MetaBase(MetaBase&&) noexcept = default;
|
| 71 |
+
MetaBase& operator=(MetaBase&&) noexcept = default;
|
| 72 |
+
virtual const Tensor& maybe_get_output(int64_t output_idx) = 0;
|
| 73 |
+
|
| 74 |
+
// Note: [set_output_*]
|
| 75 |
+
// See: https://github.com/pytorch/pytorch/issues/69813
|
| 76 |
+
// Whenever defining the output properties in the META function of a
|
| 77 |
+
// structured kernel (what was usually done with `set_output`), use one of
|
| 78 |
+
// these 3 variants, instead. In order to decide which variant to use, check
|
| 79 |
+
// the following decision tree:
|
| 80 |
+
//
|
| 81 |
+
// - Can the kernel you are going to implement support output tensors
|
| 82 |
+
// with arbitrary strides?
|
| 83 |
+
// |
|
| 84 |
+
// -- YES: `set_output_raw_strided`
|
| 85 |
+
// |
|
| 86 |
+
// -- NO: Should the output tensor strides be contiguous?
|
| 87 |
+
// |
|
| 88 |
+
// -- YES: `set_output_contiguous`
|
| 89 |
+
// |
|
| 90 |
+
// -- NO: `set_output_strided`
|
| 91 |
+
//
|
| 92 |
+
// Use this function whenever the kernel requires specific strides for the
|
| 93 |
+
// output. If `strides` does not match the given output strides, proxy outputs
|
| 94 |
+
// will be created and passed to the IMPL function.
|
| 95 |
+
virtual void set_output_strided(
|
| 96 |
+
int64_t output_idx [[maybe_unused]],
|
| 97 |
+
IntArrayRef sizes [[maybe_unused]],
|
| 98 |
+
IntArrayRef strides [[maybe_unused]],
|
| 99 |
+
TensorOptions options [[maybe_unused]],
|
| 100 |
+
DimnameList names [[maybe_unused]] = {}) {
|
| 101 |
+
TORCH_INTERNAL_ASSERT(false, "set_output_strided not implemented.");
|
| 102 |
+
}
|
| 103 |
+
|
| 104 |
+
// Use this function whenever the kernel knows how to handle arbitrary strided
|
| 105 |
+
// outputs. This function has the same behavior as the old `set_output`: it
|
| 106 |
+
// will only re-stride if the given output was resized.
|
| 107 |
+
virtual void set_output_raw_strided(
|
| 108 |
+
int64_t output_idx [[maybe_unused]],
|
| 109 |
+
IntArrayRef sizes [[maybe_unused]],
|
| 110 |
+
IntArrayRef strides_hint [[maybe_unused]],
|
| 111 |
+
TensorOptions options [[maybe_unused]],
|
| 112 |
+
DimnameList names [[maybe_unused]] = {}) {
|
| 113 |
+
TORCH_INTERNAL_ASSERT(false, "set_output_strided not implemented.");
|
| 114 |
+
}
|
| 115 |
+
|
| 116 |
+
// Use this function if the kernel requires contiguous strides.
|
| 117 |
+
// Alias for `set_output_strided`, but with contiguous strides.
|
| 118 |
+
void set_output_contiguous(
|
| 119 |
+
int64_t output_idx,
|
| 120 |
+
IntArrayRef sizes,
|
| 121 |
+
TensorOptions options,
|
| 122 |
+
DimnameList names = {}) {
|
| 123 |
+
auto strides = c10::contiguous_strides(sizes);
|
| 124 |
+
set_output_strided(output_idx, sizes, strides, options, names);
|
| 125 |
+
}
|
| 126 |
+
|
| 127 |
+
// Returns a reference to an undefined tensor if there is no presupplied
|
| 128 |
+
// output
|
| 129 |
+
const Tensor& maybe_get_output() {
|
| 130 |
+
return maybe_get_output(0);
|
| 131 |
+
}
|
| 132 |
+
virtual ~MetaBase() = default;
|
| 133 |
+
};
|
| 134 |
+
|
| 135 |
+
} // namespace impl
|
| 136 |
+
|
| 137 |
+
} // namespace at
|
lib/python3.12/site-packages/torch/include/ATen/TensorNames.h
ADDED
|
@@ -0,0 +1,75 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <ATen/WrapDimUtils.h>
|
| 4 |
+
|
| 5 |
+
namespace at::namedinference {
|
| 6 |
+
|
| 7 |
+
// TensorName and TensorNames are wrappers around Dimname and DimnameList
|
| 8 |
+
// that contain helper functions to make writing name inference rules easier.
|
| 9 |
+
//
|
| 10 |
+
// A TensorName represents a Dimname associated with some DimnameList (from a
|
| 11 |
+
// Tensor). This encapsulates all the information that is needed to check if
|
| 12 |
+
// names *match* and to *unify* names.
|
| 13 |
+
//
|
| 14 |
+
// Definition: Two names in two tensors *match* if they are equal, or if at
|
| 15 |
+
// least one of them is a wildcard that can be *refined* to the other name.
|
| 16 |
+
//
|
| 17 |
+
// Definition: unify(name, other) fails if the names do not match. Otherwise,
|
| 18 |
+
// it returns the most refined of name and other.
|
| 19 |
+
//
|
| 20 |
+
// Here is an example of checking if two names match.
|
| 21 |
+
// tensor: Tensor[A, None]
|
| 22 |
+
// other: Tensor[A]
|
| 23 |
+
//
|
| 24 |
+
// Let's say we wish to check if tensor.names[-1] matches other.names[-1].
|
| 25 |
+
// None (in tensor) cannot match A (in other) because if the None were refined
|
| 26 |
+
// to A, `tensor` would have duplicate names [A, A]. Therefore we need to check
|
| 27 |
+
// tensor.names [A, None] for the existence of A.
|
| 28 |
+
struct TORCH_API TensorName {
|
| 29 |
+
explicit TensorName(ArrayRef<Dimname> origin, int origin_idx)
|
| 30 |
+
: origin_(origin),
|
| 31 |
+
name_(origin[maybe_wrap_dim(
|
| 32 |
+
origin_idx,
|
| 33 |
+
static_cast<int64_t>(origin.size()))]),
|
| 34 |
+
origin_idx_(origin_idx) {}
|
| 35 |
+
|
| 36 |
+
// op_name is only used for error reporting.
|
| 37 |
+
const TensorName& unify(const TensorName& other, const char* op_name) const;
|
| 38 |
+
Dimname toDimname() const;
|
| 39 |
+
|
| 40 |
+
private:
|
| 41 |
+
ArrayRef<Dimname> origin_;
|
| 42 |
+
Dimname name_;
|
| 43 |
+
int origin_idx_; // A named tensor can have at most 64 dims.
|
| 44 |
+
|
| 45 |
+
TORCH_API friend std::ostream& operator<<(
|
| 46 |
+
std::ostream& out,
|
| 47 |
+
const TensorName& tensorname);
|
| 48 |
+
};
|
| 49 |
+
|
| 50 |
+
using TensorNameVec = SmallVector<TensorName, 10>;
|
| 51 |
+
|
| 52 |
+
struct TORCH_API TensorNames {
|
| 53 |
+
explicit TensorNames(ArrayRef<Dimname> names);
|
| 54 |
+
|
| 55 |
+
// Create TensorNames from names[start:end]. Each individual TensorName stores
|
| 56 |
+
// `names`, NOT names[start:end], because the original tensor's names are
|
| 57 |
+
// `names`.
|
| 58 |
+
explicit TensorNames(ArrayRef<Dimname> names, int64_t start, int64_t end);
|
| 59 |
+
|
| 60 |
+
// op_name is only used for error reporting.
|
| 61 |
+
TensorNames& unifyFromRightInplace(
|
| 62 |
+
const TensorNames& other,
|
| 63 |
+
const char* op_name = "unify");
|
| 64 |
+
void checkUnique(const char* op_name) const;
|
| 65 |
+
|
| 66 |
+
void append(TensorName name);
|
| 67 |
+
std::vector<Dimname> toDimnameVec() const;
|
| 68 |
+
|
| 69 |
+
private:
|
| 70 |
+
explicit TensorNames(TensorNameVec&& names) : names_(std::move(names)) {}
|
| 71 |
+
|
| 72 |
+
TensorNameVec names_;
|
| 73 |
+
};
|
| 74 |
+
|
| 75 |
+
} // namespace at::namedinference
|
lib/python3.12/site-packages/torch/include/ATen/TensorSubclassLikeUtils.h
ADDED
|
@@ -0,0 +1,88 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
#include <ATen/core/List.h>
|
| 3 |
+
#include <ATen/core/Tensor.h>
|
| 4 |
+
#include <c10/core/impl/TorchDispatchModeTLS.h>
|
| 5 |
+
|
| 6 |
+
#ifndef AT_PER_OPERATOR_HEADERS
|
| 7 |
+
#include <ATen/Functions.h>
|
| 8 |
+
#else
|
| 9 |
+
#include <ATen/ops/equal.h>
|
| 10 |
+
#endif
|
| 11 |
+
|
| 12 |
+
namespace at {
|
| 13 |
+
|
| 14 |
+
// Note [Tensor-subclass-like Tensors]
|
| 15 |
+
// Tensor-subclass-like is defined as:
|
| 16 |
+
// - a Tensor subclass (via __torch_dispatch__ in Python or extending
|
| 17 |
+
// TensorImpl in C++)
|
| 18 |
+
// - anything else that shares the same perils as Tensor subclasses.
|
| 19 |
+
// For example, many Tensor subclasses do not have storage and meta Tensors
|
| 20 |
+
// do not have storage either, so meta Tensors belong here.
|
| 21 |
+
//
|
| 22 |
+
// We should ensure that PyTorch internals supports Tensor-subclass-like
|
| 23 |
+
// objects. In particular, Tensor-subclass-like objects struggle with two
|
| 24 |
+
// classes of operations that are problematic for Tensor subclasses:
|
| 25 |
+
// 1. Because some Tensor subclasses do not have storage, .item() or
|
| 26 |
+
// .data_ptr() calls are not good.
|
| 27 |
+
// 2. Certain in-place operations can eliminate the typing of the Tensor
|
| 28 |
+
// subclass. For example:
|
| 29 |
+
// >>> torch.zeros(input.sizes(), grad.options()).diag().copy_(input)
|
| 30 |
+
// If input is a Tensor subclass, then the above ends up either erroring out
|
| 31 |
+
// or returning a regular non-Tensor-subclass Tensor!
|
| 32 |
+
|
| 33 |
+
constexpr auto kFunctorchWrappedTensors = DispatchKeySet(
|
| 34 |
+
{DispatchKey::FuncTorchGradWrapper,
|
| 35 |
+
DispatchKey::FuncTorchBatched,
|
| 36 |
+
DispatchKey::Functionalize});
|
| 37 |
+
|
| 38 |
+
constexpr auto kTensorSubclassLike =
|
| 39 |
+
kFunctorchWrappedTensors |
|
| 40 |
+
DispatchKeySet(
|
| 41 |
+
{// WARNING: DO NOT put combined backend component + functionality keys
|
| 42 |
+
// here, you will incorrectly always match on the functionality key
|
| 43 |
+
// no matter the backend component
|
| 44 |
+
DispatchKey::Batched,
|
| 45 |
+
DispatchKey::Sparse,
|
| 46 |
+
DispatchKey::SparseCsr,
|
| 47 |
+
DispatchKey::Python}) |
|
| 48 |
+
DispatchKeySet(BackendComponent::MetaBit);
|
| 49 |
+
|
| 50 |
+
inline bool isTensorSubclassLike(const Tensor& tensor) {
|
| 51 |
+
if (c10::impl::dispatch_mode_enabled())
|
| 52 |
+
return true;
|
| 53 |
+
auto key_set = tensor.unsafeGetTensorImpl()->key_set();
|
| 54 |
+
return !(key_set & kTensorSubclassLike).empty();
|
| 55 |
+
}
|
| 56 |
+
|
| 57 |
+
inline bool areAnyTensorSubclassLike(TensorList tensors) {
|
| 58 |
+
if (c10::impl::dispatch_mode_enabled())
|
| 59 |
+
return true;
|
| 60 |
+
return std::any_of(tensors.begin(), tensors.end(), isTensorSubclassLike);
|
| 61 |
+
}
|
| 62 |
+
|
| 63 |
+
inline bool areAnyOptionalTensorSubclassLike(
|
| 64 |
+
const c10::List<std::optional<Tensor>>& tensors) {
|
| 65 |
+
if (c10::impl::dispatch_mode_enabled())
|
| 66 |
+
return true;
|
| 67 |
+
return std::any_of(
|
| 68 |
+
tensors.begin(),
|
| 69 |
+
tensors.end(),
|
| 70 |
+
[](const std::optional<Tensor>& opt_tensor) {
|
| 71 |
+
return (
|
| 72 |
+
opt_tensor.has_value() && isTensorSubclassLike(opt_tensor.value()));
|
| 73 |
+
});
|
| 74 |
+
}
|
| 75 |
+
|
| 76 |
+
// Helper function to deal testing truthfulness of a scalar tensor
|
| 77 |
+
// in a Composite Compliant manner.
|
| 78 |
+
// NOTE: This function expects a scalar tensor of boolean dtype.
|
| 79 |
+
// Eg.
|
| 80 |
+
// Non-Composite Compliant Pattern : (t == 0).all().item<bool>()
|
| 81 |
+
// Composite Compliant Patter : is_salar_tensor_true((t == 0).all())
|
| 82 |
+
inline bool is_scalar_tensor_true(const Tensor& t) {
|
| 83 |
+
TORCH_INTERNAL_ASSERT(t.dim() == 0)
|
| 84 |
+
TORCH_INTERNAL_ASSERT(t.scalar_type() == kBool)
|
| 85 |
+
return at::equal(t, t.new_ones({}, t.options()));
|
| 86 |
+
}
|
| 87 |
+
|
| 88 |
+
} // namespace at
|
lib/python3.12/site-packages/torch/include/ATen/TensorUtils.h
ADDED
|
@@ -0,0 +1,190 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <ATen/DimVector.h>
|
| 4 |
+
#include <ATen/EmptyTensor.h>
|
| 5 |
+
#include <ATen/Tensor.h>
|
| 6 |
+
#include <ATen/TensorGeometry.h>
|
| 7 |
+
#include <ATen/Utils.h>
|
| 8 |
+
|
| 9 |
+
#include <utility>
|
| 10 |
+
|
| 11 |
+
// These functions are NOT in Utils.h, because this file has a dep on Tensor.h
|
| 12 |
+
|
| 13 |
+
#define TORCH_CHECK_TENSOR_ALL(cond, ...) \
|
| 14 |
+
TORCH_CHECK((cond)._is_all_true().item<bool>(), __VA_ARGS__);
|
| 15 |
+
|
| 16 |
+
namespace at {
|
| 17 |
+
|
| 18 |
+
// The following are utility functions for checking that arguments
|
| 19 |
+
// make sense. These are particularly useful for native functions,
|
| 20 |
+
// which do NO argument checking by default.
|
| 21 |
+
|
| 22 |
+
struct TORCH_API TensorArg {
|
| 23 |
+
// NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
|
| 24 |
+
const Tensor& tensor;
|
| 25 |
+
const char* name;
|
| 26 |
+
int pos; // 1-indexed
|
| 27 |
+
TensorArg(const Tensor& tensor, const char* name, int pos)
|
| 28 |
+
: tensor(tensor), name(name), pos(pos) {}
|
| 29 |
+
// Try to mitigate any possibility of dangling reference to temporaries.
|
| 30 |
+
// NOLINTNEXTLINE(cppcoreguidelines-rvalue-reference-param-not-moved)
|
| 31 |
+
TensorArg(Tensor&& tensor, const char* name, int pos) = delete;
|
| 32 |
+
const Tensor* operator->() const {
|
| 33 |
+
return &tensor;
|
| 34 |
+
}
|
| 35 |
+
const Tensor& operator*() const {
|
| 36 |
+
return tensor;
|
| 37 |
+
}
|
| 38 |
+
};
|
| 39 |
+
|
| 40 |
+
struct TORCH_API TensorGeometryArg {
|
| 41 |
+
TensorGeometry tensor;
|
| 42 |
+
const char* name;
|
| 43 |
+
int pos; // 1-indexed
|
| 44 |
+
/* implicit */ TensorGeometryArg(TensorArg arg)
|
| 45 |
+
: tensor(TensorGeometry{arg.tensor}), name(arg.name), pos(arg.pos) {}
|
| 46 |
+
TensorGeometryArg(TensorGeometry tensor, const char* name, int pos)
|
| 47 |
+
: tensor(std::move(tensor)), name(name), pos(pos) {}
|
| 48 |
+
const TensorGeometry* operator->() const {
|
| 49 |
+
return &tensor;
|
| 50 |
+
}
|
| 51 |
+
const TensorGeometry& operator*() const {
|
| 52 |
+
return tensor;
|
| 53 |
+
}
|
| 54 |
+
};
|
| 55 |
+
|
| 56 |
+
// A string describing which function did checks on its input
|
| 57 |
+
// arguments.
|
| 58 |
+
// TODO: Consider generalizing this into a call stack.
|
| 59 |
+
using CheckedFrom = const char*;
|
| 60 |
+
|
| 61 |
+
// The undefined convention: singular operators assume their arguments
|
| 62 |
+
// are defined, but functions which take multiple tensors will
|
| 63 |
+
// implicitly filter out undefined tensors (to make it easier to perform
|
| 64 |
+
// tests which should apply if the tensor is defined, and should not
|
| 65 |
+
// otherwise.)
|
| 66 |
+
//
|
| 67 |
+
// NB: This means that the n-ary operators take lists of TensorArg,
|
| 68 |
+
// not TensorGeometryArg, because the Tensor to TensorGeometry
|
| 69 |
+
// conversion will blow up if you have undefined tensors.
|
| 70 |
+
|
| 71 |
+
TORCH_API std::ostream& operator<<(
|
| 72 |
+
std::ostream& out,
|
| 73 |
+
const TensorGeometryArg& t);
|
| 74 |
+
TORCH_API void checkDim(
|
| 75 |
+
CheckedFrom c,
|
| 76 |
+
const Tensor& tensor,
|
| 77 |
+
const char* name,
|
| 78 |
+
int pos, // 1-indexed
|
| 79 |
+
int64_t dim);
|
| 80 |
+
TORCH_API void checkDim(CheckedFrom c, const TensorGeometryArg& t, int64_t dim);
|
| 81 |
+
// NB: this is an inclusive-exclusive range
|
| 82 |
+
TORCH_API void checkDimRange(
|
| 83 |
+
CheckedFrom c,
|
| 84 |
+
const TensorGeometryArg& t,
|
| 85 |
+
int64_t dim_start,
|
| 86 |
+
int64_t dim_end);
|
| 87 |
+
TORCH_API void checkSameDim(
|
| 88 |
+
CheckedFrom c,
|
| 89 |
+
const TensorGeometryArg& t1,
|
| 90 |
+
const TensorGeometryArg& t2);
|
| 91 |
+
TORCH_API void checkContiguous(CheckedFrom c, const TensorGeometryArg& t);
|
| 92 |
+
TORCH_API void checkAllContiguous(CheckedFrom c, at::ArrayRef<TensorArg> ts);
|
| 93 |
+
TORCH_API void checkSize(
|
| 94 |
+
CheckedFrom c,
|
| 95 |
+
const TensorGeometryArg& t,
|
| 96 |
+
IntArrayRef sizes);
|
| 97 |
+
TORCH_API void checkSize_symint(
|
| 98 |
+
CheckedFrom c,
|
| 99 |
+
const TensorGeometryArg& t,
|
| 100 |
+
c10::SymIntArrayRef sizes);
|
| 101 |
+
TORCH_API void checkSize(
|
| 102 |
+
CheckedFrom c,
|
| 103 |
+
const TensorGeometryArg& t,
|
| 104 |
+
int64_t dim,
|
| 105 |
+
int64_t size);
|
| 106 |
+
TORCH_API void checkSize_symint(
|
| 107 |
+
CheckedFrom c,
|
| 108 |
+
const TensorGeometryArg& t,
|
| 109 |
+
int64_t dim,
|
| 110 |
+
const c10::SymInt& size);
|
| 111 |
+
TORCH_API void checkNumel(
|
| 112 |
+
CheckedFrom c,
|
| 113 |
+
const TensorGeometryArg& t,
|
| 114 |
+
int64_t numel);
|
| 115 |
+
TORCH_API void checkSameNumel(
|
| 116 |
+
CheckedFrom c,
|
| 117 |
+
const TensorArg& t1,
|
| 118 |
+
const TensorArg& t2);
|
| 119 |
+
TORCH_API void checkAllSameNumel(CheckedFrom c, ArrayRef<TensorArg> tensors);
|
| 120 |
+
TORCH_API void checkScalarType(CheckedFrom c, const TensorArg& t, ScalarType s);
|
| 121 |
+
TORCH_API void checkScalarTypes(
|
| 122 |
+
CheckedFrom c,
|
| 123 |
+
const TensorArg& t,
|
| 124 |
+
at::ArrayRef<ScalarType> l);
|
| 125 |
+
TORCH_API void checkSameGPU(
|
| 126 |
+
CheckedFrom c,
|
| 127 |
+
const TensorArg& t1,
|
| 128 |
+
const TensorArg& t2);
|
| 129 |
+
TORCH_API void checkAllSameGPU(CheckedFrom c, ArrayRef<TensorArg> tensors);
|
| 130 |
+
TORCH_API void checkSameType(
|
| 131 |
+
CheckedFrom c,
|
| 132 |
+
const TensorArg& t1,
|
| 133 |
+
const TensorArg& t2);
|
| 134 |
+
TORCH_API void checkAllSameType(CheckedFrom c, ArrayRef<TensorArg> tensors);
|
| 135 |
+
TORCH_API void checkSameSize(
|
| 136 |
+
CheckedFrom c,
|
| 137 |
+
const TensorArg& t1,
|
| 138 |
+
const TensorArg& t2);
|
| 139 |
+
TORCH_API void checkAllSameSize(CheckedFrom c, ArrayRef<TensorArg> tensors);
|
| 140 |
+
TORCH_API void checkDefined(CheckedFrom c, const TensorArg& t);
|
| 141 |
+
TORCH_API void checkAllDefined(CheckedFrom c, at::ArrayRef<TensorArg> t);
|
| 142 |
+
|
| 143 |
+
// FixMe: does TensorArg slow things down?
|
| 144 |
+
TORCH_API void checkBackend(
|
| 145 |
+
CheckedFrom c,
|
| 146 |
+
at::ArrayRef<Tensor> t,
|
| 147 |
+
at::Backend backend);
|
| 148 |
+
|
| 149 |
+
TORCH_API void checkDeviceType(
|
| 150 |
+
CheckedFrom c,
|
| 151 |
+
at::ArrayRef<Tensor> tensors,
|
| 152 |
+
at::DeviceType device_type);
|
| 153 |
+
|
| 154 |
+
TORCH_API void checkLayout(CheckedFrom c, const Tensor& t, Layout layout);
|
| 155 |
+
|
| 156 |
+
TORCH_API void checkLayout(
|
| 157 |
+
CheckedFrom c,
|
| 158 |
+
at::ArrayRef<Tensor> tensors,
|
| 159 |
+
at::Layout layout);
|
| 160 |
+
|
| 161 |
+
// Methods for getting data_ptr if tensor is defined
|
| 162 |
+
TORCH_API void* maybe_data_ptr(const Tensor& tensor);
|
| 163 |
+
TORCH_API void* maybe_data_ptr(const TensorArg& tensor);
|
| 164 |
+
|
| 165 |
+
TORCH_API void check_dim_size(
|
| 166 |
+
const Tensor& tensor,
|
| 167 |
+
int64_t dim,
|
| 168 |
+
int64_t dim_size,
|
| 169 |
+
int64_t size);
|
| 170 |
+
|
| 171 |
+
namespace detail {
|
| 172 |
+
TORCH_API std::vector<int64_t> defaultStrides(IntArrayRef sizes);
|
| 173 |
+
|
| 174 |
+
TORCH_API std::optional<std::vector<int64_t>> computeStride(
|
| 175 |
+
IntArrayRef oldshape,
|
| 176 |
+
IntArrayRef oldstride,
|
| 177 |
+
IntArrayRef newshape);
|
| 178 |
+
|
| 179 |
+
TORCH_API std::optional<SymDimVector> computeStride(
|
| 180 |
+
c10::SymIntArrayRef oldshape,
|
| 181 |
+
c10::SymIntArrayRef oldstride,
|
| 182 |
+
c10::SymIntArrayRef newshape);
|
| 183 |
+
|
| 184 |
+
TORCH_API std::optional<DimVector> computeStride(
|
| 185 |
+
IntArrayRef oldshape,
|
| 186 |
+
IntArrayRef oldstride,
|
| 187 |
+
const DimVector& newshape);
|
| 188 |
+
|
| 189 |
+
} // namespace detail
|
| 190 |
+
} // namespace at
|
lib/python3.12/site-packages/torch/include/ATen/WrapDimUtilsMulti.h
ADDED
|
@@ -0,0 +1,44 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <ATen/WrapDimUtils.h>
|
| 4 |
+
#include <c10/core/TensorImpl.h>
|
| 5 |
+
#include <c10/util/irange.h>
|
| 6 |
+
#include <bitset>
|
| 7 |
+
#include <sstream>
|
| 8 |
+
|
| 9 |
+
namespace at {
|
| 10 |
+
|
| 11 |
+
// This is in an extra file to work around strange interaction of
|
| 12 |
+
// bitset on Windows with operator overloading
|
| 13 |
+
|
| 14 |
+
constexpr size_t dim_bitset_size = 64;
|
| 15 |
+
|
| 16 |
+
inline std::bitset<dim_bitset_size> dim_list_to_bitset(
|
| 17 |
+
OptionalIntArrayRef opt_dims,
|
| 18 |
+
size_t ndims) {
|
| 19 |
+
TORCH_CHECK(
|
| 20 |
+
ndims <= dim_bitset_size,
|
| 21 |
+
"only tensors with up to ",
|
| 22 |
+
dim_bitset_size,
|
| 23 |
+
" dims are supported");
|
| 24 |
+
std::bitset<dim_bitset_size> seen;
|
| 25 |
+
if (opt_dims.has_value()) {
|
| 26 |
+
auto dims = opt_dims.value();
|
| 27 |
+
for (const auto i : c10::irange(dims.size())) {
|
| 28 |
+
size_t dim = maybe_wrap_dim(dims[i], static_cast<int64_t>(ndims));
|
| 29 |
+
TORCH_CHECK(
|
| 30 |
+
!seen[dim],
|
| 31 |
+
"dim ",
|
| 32 |
+
dim,
|
| 33 |
+
" appears multiple times in the list of dims");
|
| 34 |
+
seen[dim] = true;
|
| 35 |
+
}
|
| 36 |
+
} else {
|
| 37 |
+
for (size_t dim = 0; dim < ndims; dim++) {
|
| 38 |
+
seen[dim] = true;
|
| 39 |
+
}
|
| 40 |
+
}
|
| 41 |
+
return seen;
|
| 42 |
+
}
|
| 43 |
+
|
| 44 |
+
} // namespace at
|
lib/python3.12/site-packages/torch/include/ATen/autocast_mode.h
ADDED
|
@@ -0,0 +1,911 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <ATen/ATen.h>
|
| 4 |
+
#include <ATen/NativeFunctions.h>
|
| 5 |
+
#include <ATen/Operators.h>
|
| 6 |
+
#include <torch/library.h>
|
| 7 |
+
|
| 8 |
+
#include <c10/core/impl/LocalDispatchKeySet.h>
|
| 9 |
+
#include <c10/util/intrusive_ptr.h>
|
| 10 |
+
|
| 11 |
+
namespace at::autocast {
|
| 12 |
+
|
| 13 |
+
TORCH_API bool is_autocast_enabled(at::DeviceType device_type);
|
| 14 |
+
TORCH_API void set_autocast_enabled(at::DeviceType device_type, bool enabled);
|
| 15 |
+
TORCH_API at::ScalarType get_autocast_dtype(at::DeviceType device_type);
|
| 16 |
+
TORCH_API void set_autocast_dtype(
|
| 17 |
+
at::DeviceType device_type,
|
| 18 |
+
at::ScalarType dtype);
|
| 19 |
+
TORCH_API void clear_cache();
|
| 20 |
+
TORCH_API int increment_nesting();
|
| 21 |
+
TORCH_API int decrement_nesting();
|
| 22 |
+
TORCH_API bool is_autocast_cache_enabled();
|
| 23 |
+
TORCH_API void set_autocast_cache_enabled(bool enabled);
|
| 24 |
+
|
| 25 |
+
// deprecated CUDA-specific autocast APIs
|
| 26 |
+
C10_DEPRECATED_MESSAGE(
|
| 27 |
+
"at::autocast::is_enabled() is deprecated. Please use at::autocast::is_autocast_enabled(at::kCUDA) instead.")
|
| 28 |
+
TORCH_API inline bool is_enabled() {
|
| 29 |
+
TORCH_WARN_DEPRECATION(
|
| 30 |
+
"at::autocast::",
|
| 31 |
+
__func__,
|
| 32 |
+
"() is deprecated. Please use at::autocast::is_autocast_enabled(at::kCUDA) instead.")
|
| 33 |
+
return is_autocast_enabled(at::kCUDA);
|
| 34 |
+
}
|
| 35 |
+
C10_DEPRECATED_MESSAGE(
|
| 36 |
+
"at::autocast::set_enabled(enabled) is deprecated. Please use at::autocast::set_autocast_enabled(at::kCUDA, enabled) instead.")
|
| 37 |
+
TORCH_API inline void set_enabled(bool enabled) {
|
| 38 |
+
TORCH_WARN_DEPRECATION(
|
| 39 |
+
"at::autocast::",
|
| 40 |
+
__func__,
|
| 41 |
+
"(enabled) is deprecated. Please use at::autocast::set_autocast_enabled(at::kCUDA, enabled) instead.")
|
| 42 |
+
set_autocast_enabled(at::kCUDA, enabled);
|
| 43 |
+
}
|
| 44 |
+
C10_DEPRECATED_MESSAGE(
|
| 45 |
+
"at::autocast::get_autocast_gpu_dtype() is deprecated. Please use at::autocast::get_autocast_dtype(at::kCUDA) instead.")
|
| 46 |
+
TORCH_API inline at::ScalarType get_autocast_gpu_dtype() {
|
| 47 |
+
TORCH_WARN_DEPRECATION(
|
| 48 |
+
"at::autocast::",
|
| 49 |
+
__func__,
|
| 50 |
+
"() is deprecated. Please use at::autocast::get_autocast_dtype(at::kCUDA) instead.")
|
| 51 |
+
return get_autocast_dtype(at::kCUDA);
|
| 52 |
+
}
|
| 53 |
+
C10_DEPRECATED_MESSAGE(
|
| 54 |
+
"at::autocast::set_autocast_gpu_dtype(dtype) is deprecated. Please use at::autocast::set_autocast_dtype(at::kCUDA, dtype) instead.")
|
| 55 |
+
TORCH_API inline void set_autocast_gpu_dtype(at::ScalarType dtype) {
|
| 56 |
+
TORCH_WARN_DEPRECATION(
|
| 57 |
+
"at::autocast::",
|
| 58 |
+
__func__,
|
| 59 |
+
"(dtype) is deprecated. Please use at::autocast::set_autocast_dtype(at::kCUDA, dtype) instead.")
|
| 60 |
+
set_autocast_dtype(at::kCUDA, dtype);
|
| 61 |
+
}
|
| 62 |
+
|
| 63 |
+
#define DECLARE_DEPRECATED_AUTOCAST_APIS(name, device_type) \
|
| 64 |
+
C10_DEPRECATED_MESSAGE( \
|
| 65 |
+
"at::autocast::is_" #name \
|
| 66 |
+
"_enabled() is deprecated. Please use at::autocast::is_autocast_enabled(" #device_type \
|
| 67 |
+
") instead.") \
|
| 68 |
+
TORCH_API inline bool is_##name##_enabled() { \
|
| 69 |
+
TORCH_WARN_DEPRECATION( \
|
| 70 |
+
"at::autocast::", \
|
| 71 |
+
__func__, \
|
| 72 |
+
"() is deprecated. Please use at::autocast::is_autocast_enabled(" #device_type \
|
| 73 |
+
") instead.") \
|
| 74 |
+
return is_autocast_enabled(device_type); \
|
| 75 |
+
} \
|
| 76 |
+
\
|
| 77 |
+
C10_DEPRECATED_MESSAGE( \
|
| 78 |
+
"at::autocast::set_" #name \
|
| 79 |
+
"_enabled(enabled) is deprecated. Please use at::autocast::set_autocast_enabled(" #device_type \
|
| 80 |
+
", enabled) instead.") \
|
| 81 |
+
TORCH_API inline void set_##name##_enabled(bool enabled) { \
|
| 82 |
+
TORCH_WARN_DEPRECATION( \
|
| 83 |
+
"at::autocast::", \
|
| 84 |
+
__func__, \
|
| 85 |
+
"(enabled) is deprecated. Please use at::autocast::set_autocast_enabled(" #device_type \
|
| 86 |
+
", enabled) instead.") \
|
| 87 |
+
set_autocast_enabled(device_type, enabled); \
|
| 88 |
+
} \
|
| 89 |
+
\
|
| 90 |
+
C10_DEPRECATED_MESSAGE( \
|
| 91 |
+
"at::autocast::get_autocast_" #name \
|
| 92 |
+
"_dtype() is deprecated. Please use at::autocast::get_autocast_dtype(" #device_type \
|
| 93 |
+
") instead.") \
|
| 94 |
+
TORCH_API inline at::ScalarType get_autocast_##name##_dtype() { \
|
| 95 |
+
TORCH_WARN_DEPRECATION( \
|
| 96 |
+
"at::autocast::", \
|
| 97 |
+
__func__, \
|
| 98 |
+
"() is deprecated. Please at::autocast::get_autocast_dtype(" #device_type \
|
| 99 |
+
") instead.") \
|
| 100 |
+
return get_autocast_dtype(device_type); \
|
| 101 |
+
} \
|
| 102 |
+
\
|
| 103 |
+
C10_DEPRECATED_MESSAGE( \
|
| 104 |
+
"at::autocast::set_autocast_" #name \
|
| 105 |
+
"_dtype(dtype) is deprecated. Please use at::autocast::set_autocast_dtype(" #device_type \
|
| 106 |
+
", dtype) instead.") \
|
| 107 |
+
TORCH_API inline void set_autocast_##name##_dtype(at::ScalarType dtype) { \
|
| 108 |
+
TORCH_WARN_DEPRECATION( \
|
| 109 |
+
"at::autocast::", \
|
| 110 |
+
__func__, \
|
| 111 |
+
"(dtype) is deprecated. Please use at::autocast::set_autocast_dtype(" #device_type \
|
| 112 |
+
", dtype) instead.") \
|
| 113 |
+
set_autocast_dtype(device_type, dtype); \
|
| 114 |
+
}
|
| 115 |
+
|
| 116 |
+
#define AT_FORALL_DEPRECATED_AUTOCAST_BAKCNEDS(_) \
|
| 117 |
+
_(cpu, at::kCPU) \
|
| 118 |
+
_(xpu, at::kXPU) \
|
| 119 |
+
_(xla, at::kXLA) \
|
| 120 |
+
_(hpu, at::kHPU) \
|
| 121 |
+
_(ipu, at::kIPU) \
|
| 122 |
+
_(privateuseone, at::kPrivateUse1)
|
| 123 |
+
|
| 124 |
+
// deprecated other backend specific autocast APIs
|
| 125 |
+
AT_FORALL_DEPRECATED_AUTOCAST_BAKCNEDS(DECLARE_DEPRECATED_AUTOCAST_APIS)
|
| 126 |
+
|
| 127 |
+
namespace {
|
| 128 |
+
inline bool is_autocast_eligible(
|
| 129 |
+
const Tensor& tensor,
|
| 130 |
+
c10::DeviceType device_type) {
|
| 131 |
+
switch (device_type) {
|
| 132 |
+
case c10::DeviceType::CUDA:
|
| 133 |
+
return (tensor.is_cuda() || tensor.is_xla()) &&
|
| 134 |
+
tensor.is_floating_point();
|
| 135 |
+
case c10::DeviceType::CPU:
|
| 136 |
+
return (tensor.is_cpu() || tensor.is_mkldnn()) &&
|
| 137 |
+
tensor.is_floating_point();
|
| 138 |
+
case c10::DeviceType::XPU:
|
| 139 |
+
return tensor.is_xpu() && tensor.is_floating_point();
|
| 140 |
+
case c10::DeviceType::IPU:
|
| 141 |
+
return tensor.is_ipu() && tensor.is_floating_point();
|
| 142 |
+
case c10::DeviceType::HPU:
|
| 143 |
+
return tensor.is_hpu() && tensor.is_floating_point();
|
| 144 |
+
case c10::DeviceType::XLA:
|
| 145 |
+
return tensor.is_xla() && tensor.is_floating_point();
|
| 146 |
+
case c10::DeviceType::PrivateUse1:
|
| 147 |
+
return tensor.is_privateuseone() && tensor.is_floating_point();
|
| 148 |
+
case c10::DeviceType::MPS:
|
| 149 |
+
return tensor.is_mps() && tensor.is_floating_point();
|
| 150 |
+
default:
|
| 151 |
+
return false;
|
| 152 |
+
}
|
| 153 |
+
}
|
| 154 |
+
} // namespace
|
| 155 |
+
|
| 156 |
+
inline DispatchKey get_autocast_dispatch_key_from_device_type(
|
| 157 |
+
c10::DeviceType device_type) {
|
| 158 |
+
switch (device_type) {
|
| 159 |
+
case c10::DeviceType::CUDA:
|
| 160 |
+
return DispatchKey::Autocast;
|
| 161 |
+
case c10::DeviceType::CPU:
|
| 162 |
+
return DispatchKey::AutocastCPU;
|
| 163 |
+
case c10::DeviceType::XPU:
|
| 164 |
+
return DispatchKey::AutocastXPU;
|
| 165 |
+
case c10::DeviceType::IPU:
|
| 166 |
+
return DispatchKey::AutocastIPU;
|
| 167 |
+
case c10::DeviceType::HPU:
|
| 168 |
+
return DispatchKey::AutocastHPU;
|
| 169 |
+
case c10::DeviceType::XLA:
|
| 170 |
+
return DispatchKey::AutocastXLA;
|
| 171 |
+
case c10::DeviceType::PrivateUse1:
|
| 172 |
+
return DispatchKey::AutocastPrivateUse1;
|
| 173 |
+
case c10::DeviceType::MPS:
|
| 174 |
+
return DispatchKey::AutocastMPS;
|
| 175 |
+
default:
|
| 176 |
+
throw std::runtime_error(
|
| 177 |
+
"unknown device type for autocast in get_autocast_dispatch_key_from_device_type");
|
| 178 |
+
}
|
| 179 |
+
}
|
| 180 |
+
|
| 181 |
+
inline bool is_autocast_available(c10::DeviceType device_type) {
|
| 182 |
+
if (device_type == at::kCPU || device_type == at::kCUDA ||
|
| 183 |
+
device_type == at::kXPU || device_type == at::kIPU ||
|
| 184 |
+
device_type == at::kHPU || device_type == at::kXLA ||
|
| 185 |
+
device_type == at::kPrivateUse1 || device_type == at::kMPS) {
|
| 186 |
+
return true;
|
| 187 |
+
} else {
|
| 188 |
+
return false;
|
| 189 |
+
}
|
| 190 |
+
}
|
| 191 |
+
|
| 192 |
+
inline at::ScalarType get_lower_precision_fp_from_device_type(
|
| 193 |
+
c10::DeviceType device_type) {
|
| 194 |
+
if (is_autocast_available(device_type)) {
|
| 195 |
+
return get_autocast_dtype(device_type);
|
| 196 |
+
} else {
|
| 197 |
+
throw std::runtime_error(
|
| 198 |
+
"unknown device type for autocast in get_lower_precision_fp_from_device_type");
|
| 199 |
+
}
|
| 200 |
+
}
|
| 201 |
+
|
| 202 |
+
/********************************************************************
|
| 203 |
+
Logic to extract the promote type from any Tensor or TensorList args.
|
| 204 |
+
********************************************************************/
|
| 205 |
+
|
| 206 |
+
// Overload to catch Tensor args.
|
| 207 |
+
// If nextArg is floating-point, compare its scalar_type with our
|
| 208 |
+
// current best guess for the promote type, and update if necessary.
|
| 209 |
+
inline at::ScalarType prioritize(
|
| 210 |
+
at::ScalarType current,
|
| 211 |
+
const Tensor& nextArg,
|
| 212 |
+
c10::DeviceType device_type = c10::DeviceType::CUDA) {
|
| 213 |
+
if (current == at::kDouble) {
|
| 214 |
+
TORCH_CHECK(false, "promote type is double in at::autocast::prioritize");
|
| 215 |
+
return current;
|
| 216 |
+
}
|
| 217 |
+
at::ScalarType lower_precision_fp =
|
| 218 |
+
get_lower_precision_fp_from_device_type(device_type);
|
| 219 |
+
if (is_autocast_eligible(nextArg, device_type)) {
|
| 220 |
+
auto next = nextArg.scalar_type();
|
| 221 |
+
if (next == at::kDouble) {
|
| 222 |
+
return current; // ignores double tensors
|
| 223 |
+
} else if (current == at::kFloat || next == at::kFloat) {
|
| 224 |
+
return at::kFloat; // prioritizes float over lower_precision_fp
|
| 225 |
+
} else if (current == lower_precision_fp && next == lower_precision_fp) {
|
| 226 |
+
return lower_precision_fp;
|
| 227 |
+
} else {
|
| 228 |
+
TORCH_CHECK(
|
| 229 |
+
false, "Unexpected floating ScalarType in at::autocast::prioritize");
|
| 230 |
+
return current;
|
| 231 |
+
}
|
| 232 |
+
} else {
|
| 233 |
+
return current;
|
| 234 |
+
}
|
| 235 |
+
}
|
| 236 |
+
|
| 237 |
+
// Overload to catch TensorList args (for e.g. cat, stack).
|
| 238 |
+
// Reuses the overload above to process each Tensor in the list.
|
| 239 |
+
inline at::ScalarType prioritize(
|
| 240 |
+
at::ScalarType current,
|
| 241 |
+
const TensorList& list,
|
| 242 |
+
c10::DeviceType device_type = c10::DeviceType::CUDA) {
|
| 243 |
+
for (const auto& tensor : list) {
|
| 244 |
+
current = prioritize(current, tensor, device_type);
|
| 245 |
+
}
|
| 246 |
+
return current;
|
| 247 |
+
}
|
| 248 |
+
|
| 249 |
+
inline at::ScalarType prioritize(
|
| 250 |
+
at::ScalarType current,
|
| 251 |
+
const ITensorListRef& list,
|
| 252 |
+
c10::DeviceType device_type = c10::DeviceType::CUDA) {
|
| 253 |
+
for (const auto& tensor : list) {
|
| 254 |
+
current = prioritize(current, tensor, device_type);
|
| 255 |
+
}
|
| 256 |
+
return current;
|
| 257 |
+
}
|
| 258 |
+
|
| 259 |
+
// Template to catch non-Tensor args (no-op that returns current best guess)
|
| 260 |
+
template <typename T>
|
| 261 |
+
inline at::ScalarType prioritize(
|
| 262 |
+
at::ScalarType current,
|
| 263 |
+
T nextArg,
|
| 264 |
+
c10::DeviceType device_type = c10::DeviceType::CUDA) {
|
| 265 |
+
return current;
|
| 266 |
+
}
|
| 267 |
+
|
| 268 |
+
// Overload for the tail case.
|
| 269 |
+
inline at::ScalarType promote_type(
|
| 270 |
+
at::ScalarType current,
|
| 271 |
+
c10::DeviceType device_type) {
|
| 272 |
+
return current;
|
| 273 |
+
}
|
| 274 |
+
|
| 275 |
+
// Unpack args and determine if incoming lower_precision_fp tensors need to be
|
| 276 |
+
// promoted to float32. Non-Tensor arguments are ignored.
|
| 277 |
+
template <typename Arg0, typename... Args>
|
| 278 |
+
inline at::ScalarType promote_type(
|
| 279 |
+
at::ScalarType current,
|
| 280 |
+
c10::DeviceType device_type,
|
| 281 |
+
Arg0 arg0,
|
| 282 |
+
Args... args) {
|
| 283 |
+
auto new_current = prioritize(current, arg0, device_type);
|
| 284 |
+
return promote_type(new_current, device_type, args...);
|
| 285 |
+
}
|
| 286 |
+
|
| 287 |
+
/****************************************************
|
| 288 |
+
Logic to apply cached casting to any Tensor argument.
|
| 289 |
+
****************************************************/
|
| 290 |
+
inline bool is_eligible(
|
| 291 |
+
const Tensor& arg,
|
| 292 |
+
c10::DeviceType device_type = c10::DeviceType::CUDA) {
|
| 293 |
+
return (
|
| 294 |
+
arg.defined() && is_autocast_eligible(arg, device_type) &&
|
| 295 |
+
(arg.scalar_type() != at::kDouble));
|
| 296 |
+
}
|
| 297 |
+
|
| 298 |
+
// Overload to catch Tensor args
|
| 299 |
+
TORCH_API Tensor cached_cast(
|
| 300 |
+
at::ScalarType to_type,
|
| 301 |
+
const Tensor& arg,
|
| 302 |
+
c10::DeviceType device_type = c10::DeviceType::CUDA);
|
| 303 |
+
|
| 304 |
+
// Overload to process std::optional<Tensor>
|
| 305 |
+
inline std::optional<Tensor> cached_cast(
|
| 306 |
+
at::ScalarType to_type,
|
| 307 |
+
const std::optional<Tensor>& arg,
|
| 308 |
+
c10::DeviceType device_type = c10::DeviceType::CUDA) {
|
| 309 |
+
if (arg.has_value()) {
|
| 310 |
+
return cached_cast(to_type, *arg, device_type);
|
| 311 |
+
} else {
|
| 312 |
+
return std::nullopt;
|
| 313 |
+
}
|
| 314 |
+
}
|
| 315 |
+
|
| 316 |
+
// Overload to process TensorLists
|
| 317 |
+
inline std::vector<Tensor> cached_cast(
|
| 318 |
+
at::ScalarType to_type,
|
| 319 |
+
const TensorList& arg,
|
| 320 |
+
c10::DeviceType device_type = c10::DeviceType::CUDA) {
|
| 321 |
+
std::vector<Tensor> vec;
|
| 322 |
+
vec.reserve(arg.size());
|
| 323 |
+
for (const auto& t : arg) {
|
| 324 |
+
vec.emplace_back(cached_cast(to_type, t, device_type));
|
| 325 |
+
}
|
| 326 |
+
return vec;
|
| 327 |
+
}
|
| 328 |
+
|
| 329 |
+
inline std::vector<Tensor> cached_cast(
|
| 330 |
+
at::ScalarType to_type,
|
| 331 |
+
const ITensorListRef& arg,
|
| 332 |
+
c10::DeviceType device_type = c10::DeviceType::CUDA) {
|
| 333 |
+
std::vector<Tensor> vec;
|
| 334 |
+
vec.reserve(arg.size());
|
| 335 |
+
for (const auto& t : arg) {
|
| 336 |
+
vec.emplace_back(cached_cast(to_type, t, device_type));
|
| 337 |
+
}
|
| 338 |
+
return vec;
|
| 339 |
+
}
|
| 340 |
+
|
| 341 |
+
// Template to catch non-Tensor args.
|
| 342 |
+
template <typename T>
|
| 343 |
+
inline T cached_cast(
|
| 344 |
+
at::ScalarType to_type,
|
| 345 |
+
T arg,
|
| 346 |
+
c10::DeviceType device_type = c10::DeviceType::CUDA) {
|
| 347 |
+
return arg;
|
| 348 |
+
}
|
| 349 |
+
|
| 350 |
+
/*******************************************************
|
| 351 |
+
Logic to flip an output dtype flag.
|
| 352 |
+
Keep it simple for now by assuming only one such flag is
|
| 353 |
+
present in the argument list. If I ever need a function
|
| 354 |
+
with more than flag I'll figure out something else.
|
| 355 |
+
The policy is:
|
| 356 |
+
If the user has explicity specified a dtype, respect it.
|
| 357 |
+
Otherwise, set it to the autocast type.
|
| 358 |
+
********************************************************/
|
| 359 |
+
|
| 360 |
+
// Overload to catch dtype flags
|
| 361 |
+
std::optional<ScalarType> inline set_opt_dtype(
|
| 362 |
+
at::ScalarType to_type,
|
| 363 |
+
const std::optional<ScalarType>& dtype) {
|
| 364 |
+
return dtype.has_value() ? dtype : to_type;
|
| 365 |
+
}
|
| 366 |
+
|
| 367 |
+
// Template to catch other args
|
| 368 |
+
template <typename T>
|
| 369 |
+
inline T set_opt_dtype(at::ScalarType to_type, T arg) {
|
| 370 |
+
return arg;
|
| 371 |
+
}
|
| 372 |
+
|
| 373 |
+
template <typename... Args>
|
| 374 |
+
inline bool firstarg_is_eligible(
|
| 375 |
+
c10::DeviceType device_type,
|
| 376 |
+
const Tensor& arg,
|
| 377 |
+
Args... args) {
|
| 378 |
+
return is_eligible(arg, device_type);
|
| 379 |
+
}
|
| 380 |
+
|
| 381 |
+
template <typename... Args>
|
| 382 |
+
inline at::ScalarType type_from_firstarg(
|
| 383 |
+
c10::DeviceType device_type,
|
| 384 |
+
at::ScalarType to_type,
|
| 385 |
+
const Tensor& arg,
|
| 386 |
+
Args... args) {
|
| 387 |
+
return (is_eligible(arg, device_type) ? to_type : arg.scalar_type());
|
| 388 |
+
}
|
| 389 |
+
|
| 390 |
+
// Policies correspond to op categories that need code-divergent handling.
|
| 391 |
+
// Wrapper templates below are specialized based on a policy template parameter.
|
| 392 |
+
enum class CastPolicy : uint8_t {
|
| 393 |
+
lower_precision_fp = 0, // Cast all inputs to lower_precision_fp before
|
| 394 |
+
// running the op. Currently, lower_precision_fp is
|
| 395 |
+
// fp16 for AutocastCUDA, and is defined by user
|
| 396 |
+
// (default bf16) for AutocastCPU or other device.
|
| 397 |
+
fp32, // Cast all inputs to at::kFloat before running the op.
|
| 398 |
+
fp32_set_opt_dtype, // Treats functions (like softmax) that
|
| 399 |
+
// 1. we'd like to run in fp32 and
|
| 400 |
+
// 2. have a std::optional<ScalarType> arg that controls
|
| 401 |
+
// the output type.
|
| 402 |
+
// fp32_set_opt_dtype wrappers' policy is: if the output
|
| 403 |
+
// type is already set, don't touch it, otherwise, set
|
| 404 |
+
// it to at::kFloat.
|
| 405 |
+
fp32_append_dtype, // Treats functions (like norm) that
|
| 406 |
+
// 1. we'd like to run in fp32 and
|
| 407 |
+
// 2. have some overloads that accept an output type and
|
| 408 |
+
// other overloads that don't.
|
| 409 |
+
// fp32_append_dtype wrappers wrap the overloads that don't
|
| 410 |
+
// have an output dtype.
|
| 411 |
+
// The wrapper policy is: append at::kFloat to the args,
|
| 412 |
+
// and redispatch to the type-aware overload.
|
| 413 |
+
promote, // Run in the widest dtype among several args.
|
| 414 |
+
};
|
| 415 |
+
|
| 416 |
+
/********************************************************************************************************
|
| 417 |
+
Templates to provide wrapper functions
|
| 418 |
+
|
| 419 |
+
I'm copying the pattern used in core/boxing/impl/WrapFunctionIntoFunctor.h to
|
| 420 |
+
extract args and return type. (see also
|
| 421 |
+
https://stackoverflow.com/questions/46533698/how-to-deduce-argument-list-from-function-pointer)
|
| 422 |
+
|
| 423 |
+
This strategy uses an exterior "WrapFunction" that extracts arguments on behalf
|
| 424 |
+
of (in my case several specializations of) an interior "WrapFunction_".
|
| 425 |
+
Interior WrapFunction_ specializations are defined for each CastPolicy.
|
| 426 |
+
********************************************************************************************************/
|
| 427 |
+
|
| 428 |
+
// Base template for WrapFunction_, which is specialized to contain a "call"
|
| 429 |
+
// method each CastPolicy
|
| 430 |
+
template <
|
| 431 |
+
CastPolicy policy,
|
| 432 |
+
c10::DeviceType device_type,
|
| 433 |
+
class Redispatch,
|
| 434 |
+
Redispatch* F,
|
| 435 |
+
class Ret,
|
| 436 |
+
class ArgList>
|
| 437 |
+
struct WrapFunction_ {};
|
| 438 |
+
|
| 439 |
+
// CastPolicy::lower_precision_fp General_DeviceType
|
| 440 |
+
template <
|
| 441 |
+
c10::DeviceType device_type,
|
| 442 |
+
class Redispatch,
|
| 443 |
+
Redispatch* F,
|
| 444 |
+
class Ret,
|
| 445 |
+
class... Args>
|
| 446 |
+
struct WrapFunction_<
|
| 447 |
+
CastPolicy::lower_precision_fp,
|
| 448 |
+
device_type,
|
| 449 |
+
Redispatch,
|
| 450 |
+
F,
|
| 451 |
+
Ret,
|
| 452 |
+
guts::typelist::typelist<Args...>> {
|
| 453 |
+
static Ret call(Args... args) {
|
| 454 |
+
c10::impl::ExcludeDispatchKeyGuard no_autocast(
|
| 455 |
+
get_autocast_dispatch_key_from_device_type(device_type));
|
| 456 |
+
return (*F)(cached_cast(
|
| 457 |
+
get_lower_precision_fp_from_device_type(device_type),
|
| 458 |
+
args,
|
| 459 |
+
device_type)...);
|
| 460 |
+
}
|
| 461 |
+
};
|
| 462 |
+
|
| 463 |
+
// CastPolicy::fp32 General_DeviceType
|
| 464 |
+
template <
|
| 465 |
+
c10::DeviceType device_type,
|
| 466 |
+
class Redispatch,
|
| 467 |
+
Redispatch* F,
|
| 468 |
+
class Ret,
|
| 469 |
+
class... Args>
|
| 470 |
+
struct WrapFunction_<
|
| 471 |
+
CastPolicy::fp32,
|
| 472 |
+
device_type,
|
| 473 |
+
Redispatch,
|
| 474 |
+
F,
|
| 475 |
+
Ret,
|
| 476 |
+
guts::typelist::typelist<Args...>> {
|
| 477 |
+
static Ret call(Args... args) {
|
| 478 |
+
c10::impl::ExcludeDispatchKeyGuard no_autocast(
|
| 479 |
+
get_autocast_dispatch_key_from_device_type(device_type));
|
| 480 |
+
return (*F)(cached_cast(at::kFloat, args, device_type)...);
|
| 481 |
+
}
|
| 482 |
+
};
|
| 483 |
+
|
| 484 |
+
// CastPolicy::fp32_set_opt_dtype General_DeviceType
|
| 485 |
+
template <
|
| 486 |
+
c10::DeviceType device_type,
|
| 487 |
+
class Redispatch,
|
| 488 |
+
Redispatch* F,
|
| 489 |
+
class Ret,
|
| 490 |
+
class... Args>
|
| 491 |
+
struct WrapFunction_<
|
| 492 |
+
CastPolicy::fp32_set_opt_dtype,
|
| 493 |
+
device_type,
|
| 494 |
+
Redispatch,
|
| 495 |
+
F,
|
| 496 |
+
Ret,
|
| 497 |
+
guts::typelist::typelist<Args...>> {
|
| 498 |
+
static Ret call(Args... args) {
|
| 499 |
+
c10::impl::ExcludeDispatchKeyGuard no_autocast(
|
| 500 |
+
get_autocast_dispatch_key_from_device_type(device_type));
|
| 501 |
+
if (firstarg_is_eligible(device_type, args...)) {
|
| 502 |
+
return (*F)(set_opt_dtype(at::kFloat, args)...);
|
| 503 |
+
} else {
|
| 504 |
+
// If ineligible, calls F with unaltered args. Does not set opt dtype,
|
| 505 |
+
// because setting opt dtype explicitly may interfere with internal
|
| 506 |
+
// implicit promotion decisions.
|
| 507 |
+
return (*F)(args...);
|
| 508 |
+
}
|
| 509 |
+
}
|
| 510 |
+
};
|
| 511 |
+
|
| 512 |
+
// CastPolicy::fp32_append_dtype General_DeviceType
|
| 513 |
+
template <
|
| 514 |
+
c10::DeviceType device_type,
|
| 515 |
+
class Redispatch,
|
| 516 |
+
Redispatch* F,
|
| 517 |
+
class Ret,
|
| 518 |
+
class... Args>
|
| 519 |
+
struct WrapFunction_<
|
| 520 |
+
CastPolicy::fp32_append_dtype,
|
| 521 |
+
device_type,
|
| 522 |
+
Redispatch,
|
| 523 |
+
F,
|
| 524 |
+
Ret,
|
| 525 |
+
guts::typelist::typelist<Args...>> {
|
| 526 |
+
static Ret call(Args... args) {
|
| 527 |
+
c10::impl::ExcludeDispatchKeyGuard no_autocast(
|
| 528 |
+
get_autocast_dispatch_key_from_device_type(device_type));
|
| 529 |
+
at::ScalarType out_type =
|
| 530 |
+
type_from_firstarg(device_type, at::kFloat, args...);
|
| 531 |
+
return (*F)(args..., out_type);
|
| 532 |
+
}
|
| 533 |
+
};
|
| 534 |
+
|
| 535 |
+
// CastPolicy::promote General_DeviceType
|
| 536 |
+
template <
|
| 537 |
+
c10::DeviceType device_type,
|
| 538 |
+
class Redispatch,
|
| 539 |
+
Redispatch* F,
|
| 540 |
+
class Ret,
|
| 541 |
+
class... Args>
|
| 542 |
+
struct WrapFunction_<
|
| 543 |
+
CastPolicy::promote,
|
| 544 |
+
device_type,
|
| 545 |
+
Redispatch,
|
| 546 |
+
F,
|
| 547 |
+
Ret,
|
| 548 |
+
guts::typelist::typelist<Args...>> {
|
| 549 |
+
static Ret call(Args... args) {
|
| 550 |
+
c10::impl::ExcludeDispatchKeyGuard no_autocast(
|
| 551 |
+
get_autocast_dispatch_key_from_device_type(device_type));
|
| 552 |
+
auto to_type = promote_type(
|
| 553 |
+
get_lower_precision_fp_from_device_type(device_type),
|
| 554 |
+
device_type,
|
| 555 |
+
args...);
|
| 556 |
+
return (*F)(cached_cast(to_type, args, device_type)...);
|
| 557 |
+
}
|
| 558 |
+
};
|
| 559 |
+
|
| 560 |
+
// Wrapper to infer return_type and parameter_types for WrapFunction_ (imitating
|
| 561 |
+
// core/boxing/impl/WrapFunctionIntoFunctor.h)
|
| 562 |
+
template <
|
| 563 |
+
CastPolicy policy,
|
| 564 |
+
c10::DeviceType device_type,
|
| 565 |
+
class Registered, // The signature for which we're registering. The
|
| 566 |
+
// dispatcher's calling code invokes our registered
|
| 567 |
+
// functions with arguments matching Registered, so we
|
| 568 |
+
// register WrapFunction_::call methods with a matching
|
| 569 |
+
// signature to properly field those arguments.
|
| 570 |
+
// guts::function_traits below extracts return_type and
|
| 571 |
+
// parameter_types from Registered, which WrapFunction_
|
| 572 |
+
// templates above use to declare their call methods.
|
| 573 |
+
class Redispatch, // The signature for the function we're redispatching to.
|
| 574 |
+
// In most cases this is the same as Registered, but for
|
| 575 |
+
// some ops (for example, ops where we append a dtype)
|
| 576 |
+
// it's useful to redispatch to a function with a
|
| 577 |
+
// different signature.
|
| 578 |
+
Redispatch* F> // The actual function we're redispatching to.
|
| 579 |
+
struct WrapFunction final {
|
| 580 |
+
using type = WrapFunction_<
|
| 581 |
+
policy,
|
| 582 |
+
device_type,
|
| 583 |
+
Redispatch,
|
| 584 |
+
F,
|
| 585 |
+
typename guts::function_traits<Registered>::return_type,
|
| 586 |
+
typename guts::function_traits<Registered>::parameter_types>;
|
| 587 |
+
};
|
| 588 |
+
|
| 589 |
+
/*****************************************************************************************************************
|
| 590 |
+
This section performs load-time registration for autocast wrappers.
|
| 591 |
+
|
| 592 |
+
It's debatable at what level operations should be patched. We'd like casts to
|
| 593 |
+
be autograd-exposed and precede autograd history recording, so that for
|
| 594 |
+
lower_precision_fp ops, input tensors are saved for backward in
|
| 595 |
+
lower_precision_fp rather than fp32. Saving inputs in lower_precision_fp
|
| 596 |
+
can significantly reduce a model's memory footprint.
|
| 597 |
+
|
| 598 |
+
Option 1 (strawman): Patch only at the level of explicit calls into
|
| 599 |
+
cudnn/cublas (cudnn_convolution, etc), because those are the code paths that are
|
| 600 |
+
guaranteed to use Tensor Cores, therefore they're the ones that will benefit
|
| 601 |
+
most from lower_precision_fp. Potential pitfall: convolutions (and other ops)
|
| 602 |
+
are wrapped in several layers of at::* calls. If one of those happens to record
|
| 603 |
+
autograd history, then we've lost the opportunity to save inputs in
|
| 604 |
+
lower_precision_fp.
|
| 605 |
+
|
| 606 |
+
Option 2: Patch the Python-exposed surface of calls, to make 100% sure autograd
|
| 607 |
+
history recording can't sneak in ahead of autocast. This mirrors Apex most
|
| 608 |
+
closely.
|
| 609 |
+
|
| 610 |
+
I think Option 2 is the right answer for all ops, not just convolutions. Option
|
| 611 |
+
2 is what I implement here.
|
| 612 |
+
*****************************************************************************************************************/
|
| 613 |
+
|
| 614 |
+
/********************************************************************************************************************
|
| 615 |
+
Explicit registration for out-of-place ops
|
| 616 |
+
|
| 617 |
+
The stuff below could be codegenned. Ed said
|
| 618 |
+
> you are going to have to write the function definition at some point, I
|
| 619 |
+
wouldn't try to get clever about it Therefore, for the moment, this is all
|
| 620 |
+
copy pasted in from VariableTypeEverything.cpp with appropriate substitutions.
|
| 621 |
+
********************************************************************************************************************/
|
| 622 |
+
|
| 623 |
+
} // namespace at::autocast
|
| 624 |
+
|
| 625 |
+
#define ADD_NS(RAW_OP) at::RAW_OP
|
| 626 |
+
|
| 627 |
+
#define _KERNEL_OVERLOAD_NARG_IMPL(_0, _1, _2, N, ...) N
|
| 628 |
+
#define _KERNEL_OVERLOAD_NARG(...) \
|
| 629 |
+
C10_EXPAND_MSVC_WORKAROUND(_KERNEL_OVERLOAD_NARG_IMPL(__VA_ARGS__, 2, 1))
|
| 630 |
+
|
| 631 |
+
// Common cases where registration signature matches redispatch signature
|
| 632 |
+
// (that's why SIGNATURE is repeated in the WrapFunction instantiation)
|
| 633 |
+
#define KERNEL1(DISPATCHKEY, OP, POLICY) \
|
| 634 |
+
m.impl( \
|
| 635 |
+
TORCH_SELECTIVE_NAME("aten::" #OP), \
|
| 636 |
+
&::at::autocast::WrapFunction< \
|
| 637 |
+
::at::autocast::CastPolicy::POLICY, \
|
| 638 |
+
DISPATCHKEY, \
|
| 639 |
+
decltype(ATEN_FN(OP)), \
|
| 640 |
+
decltype(ATEN_FN(OP)), \
|
| 641 |
+
&ATEN_FN(OP)>::type::call);
|
| 642 |
+
|
| 643 |
+
#define KERNEL2(DISPATCHKEY, OP, OVERLOAD, POLICY) \
|
| 644 |
+
m.impl( \
|
| 645 |
+
TORCH_SELECTIVE_NAME("aten::" #OP "." #OVERLOAD), \
|
| 646 |
+
&::at::autocast::WrapFunction< \
|
| 647 |
+
::at::autocast::CastPolicy::POLICY, \
|
| 648 |
+
DISPATCHKEY, \
|
| 649 |
+
decltype(ATEN_FN2(OP, OVERLOAD)), \
|
| 650 |
+
decltype(ATEN_FN2(OP, OVERLOAD)), \
|
| 651 |
+
&ATEN_FN2(OP, OVERLOAD)>::type::call);
|
| 652 |
+
|
| 653 |
+
#define _KERNEL_DISPATCH(DISPATCHKEY, NARG, ...) \
|
| 654 |
+
C10_CONCATENATE(KERNEL, NARG)(DISPATCHKEY, __VA_ARGS__)
|
| 655 |
+
|
| 656 |
+
#define _KERNEL_IMPL(DISPATCHKEY, ...) \
|
| 657 |
+
_KERNEL_DISPATCH(DISPATCHKEY, _KERNEL_OVERLOAD_NARG(__VA_ARGS__), __VA_ARGS__)
|
| 658 |
+
|
| 659 |
+
// It will dispatch to KERNEL1 or KERNEL2 based on its inputs.
|
| 660 |
+
#define KERNEL(DISPATCHKEY, ...) _KERNEL_IMPL(DISPATCHKEY, __VA_ARGS__)
|
| 661 |
+
|
| 662 |
+
// Less-common but still useful case: redispatching to a function
|
| 663 |
+
// with a new signature (e.g. appending a dtype)
|
| 664 |
+
#define KERNEL_DIFFERENT_REDISPATCH_SIGNATURE( \
|
| 665 |
+
DISPATCHKEY, \
|
| 666 |
+
REDISPATCH_FUNC, \
|
| 667 |
+
REGISTER_NAME, \
|
| 668 |
+
REGISTER_SIGNATURE, \
|
| 669 |
+
REDISPATCH_SIGNATURE, \
|
| 670 |
+
POLICY) \
|
| 671 |
+
m.impl( \
|
| 672 |
+
TORCH_SELECTIVE_NAME("aten::" REGISTER_NAME), \
|
| 673 |
+
&::at::autocast::WrapFunction< \
|
| 674 |
+
::at::autocast::CastPolicy::POLICY, \
|
| 675 |
+
DISPATCHKEY, \
|
| 676 |
+
REGISTER_SIGNATURE, \
|
| 677 |
+
REDISPATCH_SIGNATURE, \
|
| 678 |
+
&REDISPATCH_FUNC>::type::call);
|
| 679 |
+
|
| 680 |
+
// KERNEL_CPU/KERNEL_DIFFERENT_REDISPATCH_SIGNATURE_CPU
|
| 681 |
+
// registration (OP, POLICY) or (OP, OVERLOAD, POLICY) for AutocastCPU
|
| 682 |
+
#define KERNEL_CPU(...) KERNEL(c10::DeviceType::CPU, __VA_ARGS__)
|
| 683 |
+
|
| 684 |
+
#define KERNEL_DIFFERENT_REDISPATCH_SIGNATURE_CPU( \
|
| 685 |
+
REDISPATCH_FUNC, \
|
| 686 |
+
REGISTER_NAME, \
|
| 687 |
+
REGISTER_SIGNATURE, \
|
| 688 |
+
REDISPATCH_SIGNATURE, \
|
| 689 |
+
POLICY) \
|
| 690 |
+
KERNEL_DIFFERENT_REDISPATCH_SIGNATURE( \
|
| 691 |
+
c10::DeviceType::CPU, \
|
| 692 |
+
REDISPATCH_FUNC, \
|
| 693 |
+
REGISTER_NAME, \
|
| 694 |
+
REGISTER_SIGNATURE, \
|
| 695 |
+
REDISPATCH_SIGNATURE, \
|
| 696 |
+
POLICY)
|
| 697 |
+
|
| 698 |
+
// KERNEL_CUDA/KERNEL_DIFFERENT_REDISPATCH_SIGNATURE_CUDA
|
| 699 |
+
// registration (OP, POLICY) or (OP, OVERLOAD, POLICY) for AutocastCUDA
|
| 700 |
+
#define KERNEL_CUDA(...) KERNEL(c10::DeviceType::CUDA, __VA_ARGS__)
|
| 701 |
+
|
| 702 |
+
#define KERNEL_DIFFERENT_REDISPATCH_SIGNATURE_CUDA( \
|
| 703 |
+
REDISPATCH_FUNC, \
|
| 704 |
+
REGISTER_NAME, \
|
| 705 |
+
REGISTER_SIGNATURE, \
|
| 706 |
+
REDISPATCH_SIGNATURE, \
|
| 707 |
+
POLICY) \
|
| 708 |
+
KERNEL_DIFFERENT_REDISPATCH_SIGNATURE( \
|
| 709 |
+
c10::DeviceType::CUDA, \
|
| 710 |
+
REDISPATCH_FUNC, \
|
| 711 |
+
REGISTER_NAME, \
|
| 712 |
+
REGISTER_SIGNATURE, \
|
| 713 |
+
REDISPATCH_SIGNATURE, \
|
| 714 |
+
POLICY)
|
| 715 |
+
|
| 716 |
+
// KERNEL_XPU/KERNEL_DIFFERENT_REDISPATCH_SIGNATURE_XPU
|
| 717 |
+
// registration (OP, POLICY) or (OP, OVERLOAD, POLICY) for AutocastXPU
|
| 718 |
+
#define KERNEL_XPU(...) KERNEL(c10::DeviceType::XPU, __VA_ARGS__)
|
| 719 |
+
|
| 720 |
+
#define KERNEL_DIFFERENT_REDISPATCH_SIGNATURE_XPU( \
|
| 721 |
+
REDISPATCH_FUNC, \
|
| 722 |
+
REGISTER_NAME, \
|
| 723 |
+
REGISTER_SIGNATURE, \
|
| 724 |
+
REDISPATCH_SIGNATURE, \
|
| 725 |
+
POLICY) \
|
| 726 |
+
KERNEL_DIFFERENT_REDISPATCH_SIGNATURE( \
|
| 727 |
+
c10::DeviceType::XPU, \
|
| 728 |
+
REDISPATCH_FUNC, \
|
| 729 |
+
REGISTER_NAME, \
|
| 730 |
+
REGISTER_SIGNATURE, \
|
| 731 |
+
REDISPATCH_SIGNATURE, \
|
| 732 |
+
POLICY)
|
| 733 |
+
|
| 734 |
+
// KERNEL_PRIVATEUSEONE/KERNEL_DIFFERENT_REDISPATCH_SIGNATURE_PRIVATEUSEONE
|
| 735 |
+
// registration (OP, POLICY) or (OP, OVERLOAD, POLICY) for AutocastPrivateUse1
|
| 736 |
+
#define KERNEL_PRIVATEUSEONE(...) \
|
| 737 |
+
KERNEL(c10::DeviceType::PrivateUse1, __VA_ARGS__)
|
| 738 |
+
|
| 739 |
+
#define KERNEL_DIFFERENT_REDISPATCH_SIGNATURE_PRIVATEUSEONE( \
|
| 740 |
+
REDISPATCH_FUNC, \
|
| 741 |
+
REGISTER_NAME, \
|
| 742 |
+
REGISTER_SIGNATURE, \
|
| 743 |
+
REDISPATCH_SIGNATURE, \
|
| 744 |
+
POLICY) \
|
| 745 |
+
KERNEL_DIFFERENT_REDISPATCH_SIGNATURE( \
|
| 746 |
+
c10::DeviceType::PrivateUse1, \
|
| 747 |
+
REDISPATCH_FUNC, \
|
| 748 |
+
REGISTER_NAME, \
|
| 749 |
+
REGISTER_SIGNATURE, \
|
| 750 |
+
REDISPATCH_SIGNATURE, \
|
| 751 |
+
POLICY)
|
| 752 |
+
|
| 753 |
+
// KERNEL_MPS
|
| 754 |
+
// registration (OP, POLICY) or (OP, OVERLOAD, POLICY) for AutocastMPS
|
| 755 |
+
#define KERNEL_MPS(...) KERNEL(c10::DeviceType::MPS, __VA_ARGS__)
|
| 756 |
+
|
| 757 |
+
// Op lists for different policies.
|
| 758 |
+
// To make sure other backends can reuse the policy op list.
|
| 759 |
+
#define AT_FORALL_LOWER_PRECISION_FP(_) \
|
| 760 |
+
_(_convolution, deprecated) \
|
| 761 |
+
_(_convolution) \
|
| 762 |
+
_(conv1d) \
|
| 763 |
+
_(conv2d) \
|
| 764 |
+
_(conv3d) \
|
| 765 |
+
_(conv_tbc) \
|
| 766 |
+
_(conv_transpose1d) \
|
| 767 |
+
_(conv_transpose2d, input) \
|
| 768 |
+
_(conv_transpose3d, input) \
|
| 769 |
+
_(convolution) \
|
| 770 |
+
_(prelu) \
|
| 771 |
+
_(addmm) \
|
| 772 |
+
_(addmv) \
|
| 773 |
+
_(addr) \
|
| 774 |
+
_(matmul) \
|
| 775 |
+
_(einsum) \
|
| 776 |
+
_(mm) \
|
| 777 |
+
_(mv) \
|
| 778 |
+
_(linalg_vecdot) \
|
| 779 |
+
_(linear) \
|
| 780 |
+
_(addbmm) \
|
| 781 |
+
_(baddbmm) \
|
| 782 |
+
_(bmm) \
|
| 783 |
+
_(chain_matmul) \
|
| 784 |
+
_(linalg_multi_dot) \
|
| 785 |
+
_(_thnn_fused_lstm_cell) \
|
| 786 |
+
_(_thnn_fused_gru_cell) \
|
| 787 |
+
_(lstm_cell) \
|
| 788 |
+
_(gru_cell) \
|
| 789 |
+
_(rnn_tanh_cell) \
|
| 790 |
+
_(rnn_relu_cell) \
|
| 791 |
+
_(_scaled_dot_product_flash_attention) \
|
| 792 |
+
_(scaled_dot_product_attention)
|
| 793 |
+
|
| 794 |
+
#define AT_FORALL_FP32(_) \
|
| 795 |
+
_(acos) \
|
| 796 |
+
_(asin) \
|
| 797 |
+
_(cosh) \
|
| 798 |
+
_(erfinv) \
|
| 799 |
+
_(exp) \
|
| 800 |
+
_(expm1) \
|
| 801 |
+
_(log) \
|
| 802 |
+
_(log10) \
|
| 803 |
+
_(log2) \
|
| 804 |
+
_(log1p) \
|
| 805 |
+
_(reciprocal) \
|
| 806 |
+
_(rsqrt) \
|
| 807 |
+
_(sinh) \
|
| 808 |
+
_(tan) \
|
| 809 |
+
_(pow, Tensor_Scalar) \
|
| 810 |
+
_(pow, Tensor_Tensor) \
|
| 811 |
+
_(pow, Scalar) \
|
| 812 |
+
_(softplus) \
|
| 813 |
+
_(layer_norm) \
|
| 814 |
+
_(native_layer_norm) \
|
| 815 |
+
_(group_norm) \
|
| 816 |
+
_(frobenius_norm, dim) \
|
| 817 |
+
_(nuclear_norm) \
|
| 818 |
+
_(nuclear_norm, dim) \
|
| 819 |
+
_(cosine_similarity) \
|
| 820 |
+
_(poisson_nll_loss) \
|
| 821 |
+
_(cosine_embedding_loss) \
|
| 822 |
+
_(nll_loss) \
|
| 823 |
+
_(nll_loss2d) \
|
| 824 |
+
_(hinge_embedding_loss) \
|
| 825 |
+
_(kl_div) \
|
| 826 |
+
_(l1_loss) \
|
| 827 |
+
_(smooth_l1_loss) \
|
| 828 |
+
_(huber_loss) \
|
| 829 |
+
_(mse_loss) \
|
| 830 |
+
_(margin_ranking_loss) \
|
| 831 |
+
_(multilabel_margin_loss) \
|
| 832 |
+
_(soft_margin_loss) \
|
| 833 |
+
_(triplet_margin_loss) \
|
| 834 |
+
_(multi_margin_loss) \
|
| 835 |
+
_(binary_cross_entropy_with_logits) \
|
| 836 |
+
_(dist) \
|
| 837 |
+
_(pdist) \
|
| 838 |
+
_(cdist) \
|
| 839 |
+
_(renorm) \
|
| 840 |
+
_(logsumexp) \
|
| 841 |
+
_(upsample_nearest1d) \
|
| 842 |
+
_(_upsample_nearest_exact1d) \
|
| 843 |
+
_(upsample_nearest2d) \
|
| 844 |
+
_(_upsample_nearest_exact2d) \
|
| 845 |
+
_(upsample_nearest3d) \
|
| 846 |
+
_(_upsample_nearest_exact3d) \
|
| 847 |
+
_(upsample_linear1d) \
|
| 848 |
+
_(upsample_bilinear2d) \
|
| 849 |
+
_(_upsample_bilinear2d_aa) \
|
| 850 |
+
_(upsample_trilinear3d) \
|
| 851 |
+
_(upsample_bicubic2d) \
|
| 852 |
+
_(_upsample_bicubic2d_aa)
|
| 853 |
+
|
| 854 |
+
#define AT_FORALL_FP32_SET_OPT_DTYPE(_) \
|
| 855 |
+
_(prod) \
|
| 856 |
+
_(prod, dim_int) \
|
| 857 |
+
_(prod, dim_Dimname) \
|
| 858 |
+
_(softmax, int) \
|
| 859 |
+
_(softmax, Dimname) \
|
| 860 |
+
_(log_softmax, int) \
|
| 861 |
+
_(log_softmax, Dimname) \
|
| 862 |
+
_(cumprod) \
|
| 863 |
+
_(cumprod, dimname) \
|
| 864 |
+
_(cumsum) \
|
| 865 |
+
_(cumsum, dimname) \
|
| 866 |
+
_(linalg_vector_norm) \
|
| 867 |
+
_(linalg_matrix_norm) \
|
| 868 |
+
_(linalg_matrix_norm, str_ord) \
|
| 869 |
+
_(sum) \
|
| 870 |
+
_(sum, dim_IntList) \
|
| 871 |
+
_(sum, dim_DimnameList)
|
| 872 |
+
|
| 873 |
+
#define AT_FORALL_DIFFERENT_REDISPATCH_SIGNATURE(_) \
|
| 874 |
+
_(ADD_NS(norm), \
|
| 875 |
+
"norm.Scalar", \
|
| 876 |
+
Tensor(const Tensor&, const Scalar&), \
|
| 877 |
+
Tensor(const Tensor&, const std::optional<Scalar>&, ScalarType), \
|
| 878 |
+
fp32_append_dtype) \
|
| 879 |
+
_(ADD_NS(norm), \
|
| 880 |
+
"norm.ScalarOpt_dim", \
|
| 881 |
+
Tensor(const Tensor&, const std::optional<Scalar>&, IntArrayRef, bool), \
|
| 882 |
+
Tensor( \
|
| 883 |
+
const Tensor&, \
|
| 884 |
+
const std::optional<Scalar>&, \
|
| 885 |
+
IntArrayRef, \
|
| 886 |
+
bool, \
|
| 887 |
+
ScalarType), \
|
| 888 |
+
fp32_append_dtype) \
|
| 889 |
+
_(ADD_NS(norm), \
|
| 890 |
+
"norm.names_ScalarOpt_dim", \
|
| 891 |
+
Tensor(const Tensor&, const std::optional<Scalar>&, DimnameList, bool), \
|
| 892 |
+
Tensor( \
|
| 893 |
+
const Tensor&, \
|
| 894 |
+
const std::optional<Scalar>&, \
|
| 895 |
+
DimnameList, \
|
| 896 |
+
bool, \
|
| 897 |
+
ScalarType), \
|
| 898 |
+
fp32_append_dtype)
|
| 899 |
+
|
| 900 |
+
#define AT_FORALL_PROMOTE(_) \
|
| 901 |
+
_(addcdiv) \
|
| 902 |
+
_(addcmul) \
|
| 903 |
+
_(atan2) \
|
| 904 |
+
_(bilinear) \
|
| 905 |
+
_(cross) \
|
| 906 |
+
_(dot) \
|
| 907 |
+
_(vdot) \
|
| 908 |
+
_(grid_sampler) \
|
| 909 |
+
_(index_put) \
|
| 910 |
+
_(tensordot) \
|
| 911 |
+
_(scatter_add)
|
lib/python3.12/site-packages/torch/include/ATen/ceil_div.h
ADDED
|
@@ -0,0 +1,24 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
#include <c10/macros/Macros.h>
|
| 3 |
+
#include <type_traits>
|
| 4 |
+
|
| 5 |
+
namespace at {
|
| 6 |
+
|
| 7 |
+
/**
|
| 8 |
+
Computes ceil(a / b)
|
| 9 |
+
*/
|
| 10 |
+
template <typename T, typename = std::enable_if_t<std::is_integral_v<T>>>
|
| 11 |
+
C10_ALWAYS_INLINE C10_HOST_DEVICE T ceil_div(T a, T b) {
|
| 12 |
+
return (a + b - 1) / b;
|
| 13 |
+
}
|
| 14 |
+
|
| 15 |
+
/**
|
| 16 |
+
Computes ceil(a / b) * b; i.e., rounds up `a` to the next highest
|
| 17 |
+
multiple of b
|
| 18 |
+
*/
|
| 19 |
+
template <typename T>
|
| 20 |
+
C10_ALWAYS_INLINE C10_HOST_DEVICE T round_up(T a, T b) {
|
| 21 |
+
return ceil_div(a, b) * b;
|
| 22 |
+
}
|
| 23 |
+
|
| 24 |
+
} // namespace at
|
lib/python3.12/site-packages/torch/include/ATen/core/ATenGeneral.h
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <c10/macros/Macros.h>
|
lib/python3.12/site-packages/torch/include/ATen/core/ATenOpList.h
ADDED
|
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <c10/macros/Export.h>
|
| 4 |
+
|
| 5 |
+
namespace c10 {
|
| 6 |
+
struct OperatorName;
|
| 7 |
+
}
|
| 8 |
+
|
| 9 |
+
namespace at {
|
| 10 |
+
|
| 11 |
+
// check if an op is a custom op (i.e. did not come from native_functions.yaml)
|
| 12 |
+
TORCH_API bool is_custom_op(const c10::OperatorName& opName);
|
| 13 |
+
}
|
lib/python3.12/site-packages/torch/include/ATen/core/ATen_fwd.h
ADDED
|
@@ -0,0 +1,46 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
#include <c10/core/QScheme.h>
|
| 3 |
+
|
| 4 |
+
// Forward declarations of core ATen types used in dispatch functions
|
| 5 |
+
namespace c10 {
|
| 6 |
+
|
| 7 |
+
template<typename T>
|
| 8 |
+
class List;
|
| 9 |
+
template<typename T>
|
| 10 |
+
class IListRef;
|
| 11 |
+
class Stream;
|
| 12 |
+
class Scalar;
|
| 13 |
+
class SymInt;
|
| 14 |
+
class SymIntList;
|
| 15 |
+
struct Storage;
|
| 16 |
+
struct TensorOptions;
|
| 17 |
+
template <typename T>
|
| 18 |
+
class ArrayRef;
|
| 19 |
+
template <typename T>
|
| 20 |
+
class OptionalArrayRef;
|
| 21 |
+
|
| 22 |
+
} // namespace c10
|
| 23 |
+
|
| 24 |
+
namespace at {
|
| 25 |
+
|
| 26 |
+
class Tensor;
|
| 27 |
+
class OptionalTensorRef;
|
| 28 |
+
struct Dimname;
|
| 29 |
+
struct Generator;
|
| 30 |
+
using TensorList = c10::ArrayRef<Tensor>;
|
| 31 |
+
using ITensorListRef = c10::IListRef<Tensor>;
|
| 32 |
+
using IOptTensorListRef = c10::IListRef<OptionalTensorRef>;
|
| 33 |
+
using DimnameList = c10::ArrayRef<Dimname>;
|
| 34 |
+
using IntArrayRef = c10::ArrayRef<int64_t>;
|
| 35 |
+
using OptionalIntArrayRef = c10::OptionalArrayRef<int64_t>;
|
| 36 |
+
using OptionalSymIntArrayRef = c10::OptionalArrayRef<c10::SymInt>;
|
| 37 |
+
|
| 38 |
+
using c10::Stream;
|
| 39 |
+
using c10::Storage;
|
| 40 |
+
using c10::QScheme;
|
| 41 |
+
using c10::Scalar;
|
| 42 |
+
using c10::SymInt;
|
| 43 |
+
using c10::SymIntList;
|
| 44 |
+
using c10::TensorOptions;
|
| 45 |
+
|
| 46 |
+
} // namespace at
|
lib/python3.12/site-packages/torch/include/ATen/core/ATen_pch.h
ADDED
|
@@ -0,0 +1,165 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
// This global header must not depend on native_functions.yaml or
|
| 2 |
+
// incremental builds will be next to useless
|
| 3 |
+
#pragma push_macro("TORCH_ASSERT_NO_OPERATORS")
|
| 4 |
+
#define TORCH_ASSERT_NO_OPERATORS
|
| 5 |
+
|
| 6 |
+
// This macro doesn't work if defined after the first time inttypes.h
|
| 7 |
+
// is included, so won't work anywhere if not defined here.
|
| 8 |
+
#ifndef __STDC_FORMAT_MACROS
|
| 9 |
+
#define __STDC_FORMAT_MACROS
|
| 10 |
+
#endif
|
| 11 |
+
#include <cinttypes>
|
| 12 |
+
|
| 13 |
+
// This list of headers was generated using a script that finds
|
| 14 |
+
// high-impact headers and then manually tweaked to remove OS specific
|
| 15 |
+
// or duplicate headers (e.g. <cassert> and <assert.h>) and to remove
|
| 16 |
+
// "impl" headers (e.g BFloat16-inl.h or complex_math.h in c10).
|
| 17 |
+
|
| 18 |
+
// To generate the initial list:
|
| 19 |
+
// 1. Build pytorch from scratch with all build caching disabled
|
| 20 |
+
// 2. Generate a build trace with ninjatracing (https://github.com/nico/ninjatracing)
|
| 21 |
+
// $ ninjatracing /path/to/pytorch/build/.ninja_log > trace_all.json
|
| 22 |
+
// 3. Run pch_gen.py from https://github.com/peterbell10/build_analysis/
|
| 23 |
+
// $ python pch_gen.py --threshold .80 --target torch_cpu --build_dir /path/to/pytorch/build --trace trace_all.json
|
| 24 |
+
// Where the threshold can be tweaked until c10 and some of ATen
|
| 25 |
+
// core are included but TORCH_ASSERT_NO_OPERATORS still passes.
|
| 26 |
+
|
| 27 |
+
#include <cerrno>
|
| 28 |
+
#include <cmath>
|
| 29 |
+
#include <cstddef>
|
| 30 |
+
#include <cstdint>
|
| 31 |
+
#include <cstdlib>
|
| 32 |
+
#include <cstring>
|
| 33 |
+
|
| 34 |
+
#include <algorithm>
|
| 35 |
+
#include <array>
|
| 36 |
+
#include <atomic>
|
| 37 |
+
#include <chrono>
|
| 38 |
+
#include <complex>
|
| 39 |
+
#include <deque>
|
| 40 |
+
#include <exception>
|
| 41 |
+
#include <functional>
|
| 42 |
+
#include <initializer_list>
|
| 43 |
+
#include <iomanip>
|
| 44 |
+
#include <iosfwd>
|
| 45 |
+
#include <iterator>
|
| 46 |
+
#include <limits>
|
| 47 |
+
#include <list>
|
| 48 |
+
#include <map>
|
| 49 |
+
#include <memory>
|
| 50 |
+
#include <mutex>
|
| 51 |
+
#include <new>
|
| 52 |
+
#include <numeric>
|
| 53 |
+
#include <ostream>
|
| 54 |
+
#include <sstream>
|
| 55 |
+
#include <stdexcept>
|
| 56 |
+
#include <string>
|
| 57 |
+
#include <tuple>
|
| 58 |
+
#include <type_traits>
|
| 59 |
+
#include <typeindex>
|
| 60 |
+
#include <typeinfo>
|
| 61 |
+
#include <unordered_map>
|
| 62 |
+
#include <unordered_set>
|
| 63 |
+
#include <utility>
|
| 64 |
+
#include <vector>
|
| 65 |
+
|
| 66 |
+
#include <c10/core/Allocator.h>
|
| 67 |
+
#include <c10/core/AutogradState.h>
|
| 68 |
+
#include <c10/core/Backend.h>
|
| 69 |
+
#include <c10/core/DefaultDtype.h>
|
| 70 |
+
#include <c10/core/Device.h>
|
| 71 |
+
#include <c10/core/DeviceType.h>
|
| 72 |
+
#include <c10/core/DispatchKey.h>
|
| 73 |
+
#include <c10/core/DispatchKeySet.h>
|
| 74 |
+
#include <c10/core/GeneratorImpl.h>
|
| 75 |
+
#include <c10/core/InferenceMode.h>
|
| 76 |
+
#include <c10/core/Layout.h>
|
| 77 |
+
#include <c10/core/MemoryFormat.h>
|
| 78 |
+
#include <c10/core/OptionalRef.h>
|
| 79 |
+
#include <c10/core/QScheme.h>
|
| 80 |
+
#include <c10/core/Scalar.h>
|
| 81 |
+
#include <c10/core/ScalarType.h>
|
| 82 |
+
#include <c10/core/ScalarTypeToTypeMeta.h>
|
| 83 |
+
#include <c10/core/Storage.h>
|
| 84 |
+
#include <c10/core/StorageImpl.h>
|
| 85 |
+
#include <c10/core/SymBool.h>
|
| 86 |
+
#include <c10/core/SymFloat.h>
|
| 87 |
+
#include <c10/core/SymInt.h>
|
| 88 |
+
#include <c10/core/SymIntArrayRef.h>
|
| 89 |
+
#include <c10/core/SymNodeImpl.h>
|
| 90 |
+
#include <c10/core/TensorImpl.h>
|
| 91 |
+
#include <c10/core/TensorOptions.h>
|
| 92 |
+
#include <c10/core/UndefinedTensorImpl.h>
|
| 93 |
+
#include <c10/core/WrapDimMinimal.h>
|
| 94 |
+
#include <c10/core/impl/LocalDispatchKeySet.h>
|
| 95 |
+
#include <c10/core/impl/PyInterpreter.h>
|
| 96 |
+
#include <c10/core/impl/SizesAndStrides.h>
|
| 97 |
+
|
| 98 |
+
#include <c10/macros/Export.h>
|
| 99 |
+
#include <c10/macros/Macros.h>
|
| 100 |
+
|
| 101 |
+
#include <c10/util/AlignOf.h>
|
| 102 |
+
#include <c10/util/ArrayRef.h>
|
| 103 |
+
#include <c10/util/BFloat16.h>
|
| 104 |
+
#include <c10/util/C++17.h>
|
| 105 |
+
#include <c10/util/ConstexprCrc.h>
|
| 106 |
+
#include <c10/util/Deprecated.h>
|
| 107 |
+
#include <c10/util/DimVector.h>
|
| 108 |
+
#include <c10/util/Exception.h>
|
| 109 |
+
#include <c10/util/ExclusivelyOwned.h>
|
| 110 |
+
#include <c10/util/Flags.h>
|
| 111 |
+
#include <c10/util/Float8_e4m3fn.h>
|
| 112 |
+
#include <c10/util/Float8_e5m2.h>
|
| 113 |
+
#include <c10/util/Float8_e4m3fnuz.h>
|
| 114 |
+
#include <c10/util/Float8_e5m2fnuz.h>
|
| 115 |
+
#include <c10/util/FunctionRef.h>
|
| 116 |
+
#include <c10/util/Half.h>
|
| 117 |
+
#include <c10/util/IdWrapper.h>
|
| 118 |
+
#include <c10/util/Logging.h>
|
| 119 |
+
#include <c10/util/MaybeOwned.h>
|
| 120 |
+
#include <c10/util/Metaprogramming.h>
|
| 121 |
+
#include <c10/util/Optional.h>
|
| 122 |
+
#include <c10/util/Registry.h>
|
| 123 |
+
#include <c10/util/SmallVector.h>
|
| 124 |
+
#include <c10/util/StringUtil.h>
|
| 125 |
+
#include <c10/util/ThreadLocalDebugInfo.h>
|
| 126 |
+
#include <c10/util/Type.h>
|
| 127 |
+
#include <c10/util/TypeCast.h>
|
| 128 |
+
#include <c10/util/TypeIndex.h>
|
| 129 |
+
#include <c10/util/TypeList.h>
|
| 130 |
+
#include <c10/util/TypeSafeSignMath.h>
|
| 131 |
+
#include <c10/util/TypeTraits.h>
|
| 132 |
+
#include <c10/util/UniqueVoidPtr.h>
|
| 133 |
+
#include <c10/util/accumulate.h>
|
| 134 |
+
#include <c10/util/bit_cast.h>
|
| 135 |
+
#include <c10/util/bits.h>
|
| 136 |
+
#include <c10/util/complex.h>
|
| 137 |
+
#include <c10/util/floating_point_utils.h>
|
| 138 |
+
#include <c10/util/intrusive_ptr.h>
|
| 139 |
+
#include <c10/util/irange.h>
|
| 140 |
+
#include <c10/util/llvmMathExtras.h>
|
| 141 |
+
#include <c10/util/python_stub.h>
|
| 142 |
+
#include <c10/util/qint32.h>
|
| 143 |
+
#include <c10/util/qint8.h>
|
| 144 |
+
#include <c10/util/quint2x4.h>
|
| 145 |
+
#include <c10/util/quint4x2.h>
|
| 146 |
+
#include <c10/util/quint8.h>
|
| 147 |
+
#include <c10/util/safe_numerics.h>
|
| 148 |
+
#include <c10/util/string_utils.h>
|
| 149 |
+
#include <c10/util/string_view.h>
|
| 150 |
+
#include <c10/util/typeid.h>
|
| 151 |
+
|
| 152 |
+
#include <ATen/StorageUtils.h>
|
| 153 |
+
#include <ATen/core/ATen_fwd.h>
|
| 154 |
+
#include <ATen/core/DeprecatedTypeProperties.h>
|
| 155 |
+
#include <ATen/core/DeprecatedTypePropertiesRegistry.h>
|
| 156 |
+
#include <ATen/core/DimVector.h>
|
| 157 |
+
#include <ATen/core/Dimname.h>
|
| 158 |
+
#include <ATen/core/Generator.h>
|
| 159 |
+
#include <ATen/core/NamedTensor.h>
|
| 160 |
+
#include <ATen/core/QuantizerBase.h>
|
| 161 |
+
#include <ATen/core/TensorAccessor.h>
|
| 162 |
+
#include <ATen/core/TensorBase.h>
|
| 163 |
+
#include <ATen/core/symbol.h>
|
| 164 |
+
|
| 165 |
+
#pragma pop_macro("TORCH_ASSERT_NO_OPERATORS")
|
lib/python3.12/site-packages/torch/include/ATen/core/Array.h
ADDED
|
@@ -0,0 +1,48 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
// A fixed-size array type usable from both host and
|
| 4 |
+
// device code.
|
| 5 |
+
|
| 6 |
+
#include <c10/macros/Macros.h>
|
| 7 |
+
#include <c10/util/irange.h>
|
| 8 |
+
|
| 9 |
+
namespace at::detail {
|
| 10 |
+
|
| 11 |
+
template <typename T, int size_>
|
| 12 |
+
struct Array {
|
| 13 |
+
// NOLINTNEXTLINE(*c-array*)
|
| 14 |
+
T data[size_];
|
| 15 |
+
|
| 16 |
+
C10_HOST_DEVICE T operator[](int i) const {
|
| 17 |
+
return data[i];
|
| 18 |
+
}
|
| 19 |
+
C10_HOST_DEVICE T& operator[](int i) {
|
| 20 |
+
return data[i];
|
| 21 |
+
}
|
| 22 |
+
#if defined(USE_ROCM)
|
| 23 |
+
C10_HOST_DEVICE Array() = default;
|
| 24 |
+
C10_HOST_DEVICE Array(const Array&) = default;
|
| 25 |
+
C10_HOST_DEVICE Array& operator=(const Array&) = default;
|
| 26 |
+
C10_HOST_DEVICE Array(Array&&) = default;
|
| 27 |
+
C10_HOST_DEVICE Array& operator=(Array&&) = default;
|
| 28 |
+
C10_HOST_DEVICE ~Array() = default;
|
| 29 |
+
#else
|
| 30 |
+
Array() = default;
|
| 31 |
+
Array(const Array&) = default;
|
| 32 |
+
Array& operator=(const Array&) = default;
|
| 33 |
+
Array(Array&&) noexcept = default;
|
| 34 |
+
Array& operator=(Array&&) noexcept = default;
|
| 35 |
+
~Array() = default;
|
| 36 |
+
#endif
|
| 37 |
+
static constexpr int size() {
|
| 38 |
+
return size_;
|
| 39 |
+
}
|
| 40 |
+
// Fill the array with x.
|
| 41 |
+
C10_HOST_DEVICE Array(T x) {
|
| 42 |
+
for (int i = 0; i < size_; i++) {
|
| 43 |
+
data[i] = x;
|
| 44 |
+
}
|
| 45 |
+
}
|
| 46 |
+
};
|
| 47 |
+
|
| 48 |
+
} // namespace at::detail
|
lib/python3.12/site-packages/torch/include/ATen/core/Backtrace.h
ADDED
|
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#include <c10/util/Backtrace.h>
|
| 2 |
+
#include <c10/util/Type.h>
|
lib/python3.12/site-packages/torch/include/ATen/core/CachingHostAllocator.h
ADDED
|
@@ -0,0 +1,465 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#include <c10/core/Allocator.h>
|
| 2 |
+
#include <c10/core/thread_pool.h>
|
| 3 |
+
#include <c10/util/CallOnce.h>
|
| 4 |
+
#include <c10/util/flat_hash_map.h>
|
| 5 |
+
#include <c10/util/llvmMathExtras.h>
|
| 6 |
+
#include <optional>
|
| 7 |
+
|
| 8 |
+
#include <deque>
|
| 9 |
+
#include <mutex>
|
| 10 |
+
#include <set>
|
| 11 |
+
|
| 12 |
+
C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wunused-parameter")
|
| 13 |
+
namespace at {
|
| 14 |
+
|
| 15 |
+
/**
|
| 16 |
+
* HostBlock is typically a fundamental memory block used in pinned memory. It
|
| 17 |
+
* is likely related to Event and Stream of device runtime. It is probably a
|
| 18 |
+
* base struct or interface that can be inherited and extended by each backend.
|
| 19 |
+
*/
|
| 20 |
+
template <typename S>
|
| 21 |
+
struct HostBlock {
|
| 22 |
+
// constructor for search key
|
| 23 |
+
HostBlock(size_t size) : size_(size) {}
|
| 24 |
+
|
| 25 |
+
HostBlock(size_t size, void* ptr) : size_(size), ptr_(ptr) {}
|
| 26 |
+
|
| 27 |
+
std::mutex mutex_;
|
| 28 |
+
size_t size_{0}; // block size in bytes
|
| 29 |
+
void* ptr_{nullptr}; // memory address
|
| 30 |
+
bool allocated_{false}; // in-use flag
|
| 31 |
+
size_t event_count_{0}; // number of related events
|
| 32 |
+
ska::flat_hash_set<S> streams_; // streams on which the block was used
|
| 33 |
+
};
|
| 34 |
+
|
| 35 |
+
template <typename B>
|
| 36 |
+
struct alignas(64) FreeBlockList {
|
| 37 |
+
std::mutex mutex_;
|
| 38 |
+
std::deque<B*> list_;
|
| 39 |
+
};
|
| 40 |
+
|
| 41 |
+
namespace {
|
| 42 |
+
// Max cached block sizes: (1 << MAX_SIZE_INDEX) bytes
|
| 43 |
+
// NOLINTNEXTLINE(misc-definitions-in-headers)
|
| 44 |
+
constexpr size_t MAX_SIZE_INDEX = 64;
|
| 45 |
+
}
|
| 46 |
+
|
| 47 |
+
/**
|
| 48 |
+
* Note [HostAllocator design]
|
| 49 |
+
* ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
| 50 |
+
* We have three key data structures - the free list which stores blocks that
|
| 51 |
+
* are not currently used, the block list which stores all blocks that have been
|
| 52 |
+
* allocated, and the event queue which stores runtime events and their
|
| 53 |
+
* corresponding blocks.
|
| 54 |
+
*
|
| 55 |
+
* Each of these are protected by a separate mutex. The key design principles
|
| 56 |
+
* are to 1) only hold each mutex for the minimal amount of time possible, 2)
|
| 57 |
+
* never do any possible expensive operations (such as CUDA runtime API calls)
|
| 58 |
+
* while holding the lock.
|
| 59 |
+
*
|
| 60 |
+
* There are four public methods: allocate, free, record_event and empty_cache.
|
| 61 |
+
* 1) In the allocate path, we first check to see if we can service our
|
| 62 |
+
* request from this free list, and otherwise we create a new block with
|
| 63 |
+
* allocate_host_memory.
|
| 64 |
+
* 2) In the free path, we insert events (if required) into the event queue,
|
| 65 |
+
* and if possible insert our block back into the free list. In allocate, we
|
| 66 |
+
* first eagerly query events until we find one that is not ready, and insert
|
| 67 |
+
* the corresponding block onto the free list if all the events recorded for a
|
| 68 |
+
* block are ready.
|
| 69 |
+
* 3) In the record_event path, we simply insert the given stream into the set
|
| 70 |
+
* of streams tracked by the specified block. This set of streams is then
|
| 71 |
+
* consumed in the free path.
|
| 72 |
+
* 4) In the empty_cache path, we flush any available blocks into the free
|
| 73 |
+
* list. Remove all element of free list, then remove them from block list and
|
| 74 |
+
* release the associated pinned memory allocation via free_block.
|
| 75 |
+
*
|
| 76 |
+
* We generalize the caching host allocator into two parts: interface and
|
| 77 |
+
* implementation. For any new backend looking to integrate with host allocator
|
| 78 |
+
* and reuse caching mechanism, these two parts are necessary to be specialized.
|
| 79 |
+
*
|
| 80 |
+
* For the implementation, we provide a CachingHostAllocatorImpl struct
|
| 81 |
+
* to abstract the caching mechanism. Any backend needs to provide a customized
|
| 82 |
+
* implementation by specializing its own public functions and the related
|
| 83 |
+
* runtime functions. Its template parameter S represents runtime Stream, E
|
| 84 |
+
* denotes runtime Event, B indicates the fundamental memory block.
|
| 85 |
+
*
|
| 86 |
+
* For the interface, we provide a CachingHostAllocatorInterface struct as an
|
| 87 |
+
* interface. Any backend needs to derive its own host allocator from this
|
| 88 |
+
* interface. Its template parameter T refers to an implementation that
|
| 89 |
+
* inherited from CachingHostAllocatorImpl.
|
| 90 |
+
*
|
| 91 |
+
* So this design can share the caching mechanism across each backend, and
|
| 92 |
+
* provide flexibility to each backend. A backend can choose to follow this
|
| 93 |
+
* implementation or reuse them by extending and overriding them as necessary.
|
| 94 |
+
* Taking CUDA as an example, it specializes runtime related functions to reuse
|
| 95 |
+
* the caching mechanism. Additionally, it extends the allocator's functionality
|
| 96 |
+
* by adding the allocWithCudaHostRegister function to support page-locking the
|
| 97 |
+
* memory range used by CUDA. Of course, you can also refer to
|
| 98 |
+
* XPUCachingHostAllocator, which is a host caching allocator supported on XPU
|
| 99 |
+
* backend, to implement a basic host caching allocator.
|
| 100 |
+
*
|
| 101 |
+
* Some of the invariants here are less strict than they could be - for example,
|
| 102 |
+
* we do not enforce that free(Block* block) => block->event_count == 0. This is
|
| 103 |
+
* for compatibility reasons, and we can explore enforcing these in subsequent
|
| 104 |
+
* versions.
|
| 105 |
+
*
|
| 106 |
+
* Note that this caching host allocator does not split larger allocations into
|
| 107 |
+
* smaller blocks, unlike the caching device allocator.
|
| 108 |
+
*/
|
| 109 |
+
|
| 110 |
+
template <
|
| 111 |
+
typename S,
|
| 112 |
+
typename E,
|
| 113 |
+
typename B = HostBlock<S>>
|
| 114 |
+
struct CachingHostAllocatorImpl {
|
| 115 |
+
virtual ~CachingHostAllocatorImpl() = default;
|
| 116 |
+
|
| 117 |
+
public:
|
| 118 |
+
// return data_ptr and block pair.
|
| 119 |
+
virtual std::pair<void*, void*> allocate(size_t size) {
|
| 120 |
+
if (size == 0) {
|
| 121 |
+
return {nullptr, nullptr};
|
| 122 |
+
}
|
| 123 |
+
|
| 124 |
+
// If we are using background threads, we can process events in the
|
| 125 |
+
// background.
|
| 126 |
+
if (!pinned_use_background_threads()) {
|
| 127 |
+
process_events();
|
| 128 |
+
}
|
| 129 |
+
|
| 130 |
+
// Round up the allocation to the nearest power of two to improve reuse.
|
| 131 |
+
// These power of two sizes are also used to index into the free list.
|
| 132 |
+
size_t roundSize = c10::llvm::PowerOf2Ceil(size);
|
| 133 |
+
|
| 134 |
+
// First, try to allocate from the free list
|
| 135 |
+
auto* block = get_free_block(roundSize);
|
| 136 |
+
if (block) {
|
| 137 |
+
return {block->ptr_, reinterpret_cast<void*>(block)};
|
| 138 |
+
}
|
| 139 |
+
|
| 140 |
+
// Check in the recently freed blocks with pending events to see if we
|
| 141 |
+
// can reuse them. Call get_free_block again after processing events
|
| 142 |
+
if (pinned_use_background_threads()) {
|
| 143 |
+
process_events_for_specific_size(roundSize);
|
| 144 |
+
block = get_free_block(roundSize);
|
| 145 |
+
if (block) {
|
| 146 |
+
return {block->ptr_, reinterpret_cast<void*>(block)};
|
| 147 |
+
}
|
| 148 |
+
|
| 149 |
+
// Launch the background thread and process events in a loop.
|
| 150 |
+
static c10::once_flag background_thread_flag;
|
| 151 |
+
c10::call_once(background_thread_flag, [this] {
|
| 152 |
+
getBackgroundThreadPool()->run([&]() {
|
| 153 |
+
while (true) {
|
| 154 |
+
process_events();
|
| 155 |
+
std::this_thread::sleep_for(std::chrono::microseconds(100));
|
| 156 |
+
}
|
| 157 |
+
});
|
| 158 |
+
});
|
| 159 |
+
}
|
| 160 |
+
|
| 161 |
+
// Slow path: if we can't allocate from the cached free list, we need
|
| 162 |
+
// to create a new block.
|
| 163 |
+
void* ptr = nullptr;
|
| 164 |
+
allocate_host_memory(roundSize, &ptr);
|
| 165 |
+
|
| 166 |
+
// Then, create a new block.
|
| 167 |
+
block = new B(roundSize, ptr);
|
| 168 |
+
block->allocated_ = true;
|
| 169 |
+
|
| 170 |
+
add_allocated_block(block);
|
| 171 |
+
return {block->ptr_, reinterpret_cast<void*>(block)};
|
| 172 |
+
}
|
| 173 |
+
|
| 174 |
+
virtual void free(void* ctx) {
|
| 175 |
+
if (!ctx) {
|
| 176 |
+
return;
|
| 177 |
+
}
|
| 178 |
+
|
| 179 |
+
// Note: we can assume that free is correctly paired with alloc, and thus we
|
| 180 |
+
// do not need to look up the ctx in blocks_.
|
| 181 |
+
auto* block = reinterpret_cast<B*>(ctx);
|
| 182 |
+
|
| 183 |
+
std::optional<std::vector<E>> events;
|
| 184 |
+
{
|
| 185 |
+
std::lock_guard<std::mutex> g(block->mutex_);
|
| 186 |
+
block->allocated_ = false;
|
| 187 |
+
if (block->streams_.empty()) {
|
| 188 |
+
TORCH_INTERNAL_ASSERT(block->event_count_ == 0);
|
| 189 |
+
} else {
|
| 190 |
+
events = std::vector<E>();
|
| 191 |
+
events->reserve(block->streams_.size());
|
| 192 |
+
for (auto stream : block->streams_) {
|
| 193 |
+
record_stream(events, stream);
|
| 194 |
+
}
|
| 195 |
+
block->event_count_ += events->size();
|
| 196 |
+
block->streams_.clear();
|
| 197 |
+
}
|
| 198 |
+
}
|
| 199 |
+
|
| 200 |
+
if (!events) {
|
| 201 |
+
auto index = size_index(block->size_);
|
| 202 |
+
std::lock_guard<std::mutex> g(free_list_[index].mutex_);
|
| 203 |
+
free_list_[index].list_.push_back(block);
|
| 204 |
+
} else {
|
| 205 |
+
// restore these events that record by used streams.
|
| 206 |
+
std::lock_guard<std::mutex> g(events_mutex_);
|
| 207 |
+
for (auto&& event : *events) {
|
| 208 |
+
events_.emplace_front(std::move(event), block);
|
| 209 |
+
}
|
| 210 |
+
}
|
| 211 |
+
}
|
| 212 |
+
|
| 213 |
+
virtual bool record_event(void* ptr, void* ctx, S stream) {
|
| 214 |
+
auto* block = reinterpret_cast<B*>(ctx);
|
| 215 |
+
|
| 216 |
+
// Note: we need to check if the passed-in `ctx` is valid. This is because
|
| 217 |
+
// `record_event` (via `CachingHostAllocator_recordEvent`) can be invoked on
|
| 218 |
+
// an arbitrary tensor, and is not guaranteed to correspond to a pinned
|
| 219 |
+
// memory allocation. Therefore, we need to check that `ctx` is valid before
|
| 220 |
+
// proceeding.
|
| 221 |
+
{
|
| 222 |
+
std::lock_guard<std::mutex> g(blocks_mutex_);
|
| 223 |
+
if (blocks_.find(block) != blocks_.end()) {
|
| 224 |
+
// Now we know this object is safe to access.
|
| 225 |
+
std::lock_guard<std::mutex> gb(block->mutex_);
|
| 226 |
+
TORCH_INTERNAL_ASSERT(block->allocated_);
|
| 227 |
+
block->streams_.insert(stream);
|
| 228 |
+
return true;
|
| 229 |
+
}
|
| 230 |
+
auto it = ptr_to_block_.find(ptr);
|
| 231 |
+
if (it != ptr_to_block_.end()) {
|
| 232 |
+
block = it->second;
|
| 233 |
+
std::lock_guard<std::mutex> g(block->mutex_);
|
| 234 |
+
TORCH_INTERNAL_ASSERT(block->allocated_);
|
| 235 |
+
block->streams_.insert(stream);
|
| 236 |
+
return true;
|
| 237 |
+
}
|
| 238 |
+
}
|
| 239 |
+
|
| 240 |
+
return false;
|
| 241 |
+
}
|
| 242 |
+
|
| 243 |
+
virtual void empty_cache() {
|
| 244 |
+
// Flush any available blocks into the free_list.
|
| 245 |
+
process_events();
|
| 246 |
+
|
| 247 |
+
// Remove all elements from the free list, remove them from the blocks
|
| 248 |
+
// list, and free the associated pinned memory allocation. This requires
|
| 249 |
+
// concurrently holding both the free list mutexes and the blocks mutex, and
|
| 250 |
+
// is the only function that concurrently holds multiple mutexes.
|
| 251 |
+
for (size_t i = 0; i < free_list_.size(); ++i) {
|
| 252 |
+
std::lock(free_list_[i].mutex_, blocks_mutex_);
|
| 253 |
+
std::lock_guard<std::mutex> gf(free_list_[i].mutex_, std::adopt_lock);
|
| 254 |
+
std::lock_guard<std::mutex> gb(blocks_mutex_, std::adopt_lock);
|
| 255 |
+
|
| 256 |
+
std::vector<B*> blocks_to_remove(free_list_[i].list_.begin(), free_list_[i].list_.end());
|
| 257 |
+
free_list_[i].list_.clear();
|
| 258 |
+
for (auto* block : blocks_to_remove) {
|
| 259 |
+
blocks_.erase(block);
|
| 260 |
+
ptr_to_block_.erase(block->ptr_);
|
| 261 |
+
free_block(block);
|
| 262 |
+
delete block;
|
| 263 |
+
}
|
| 264 |
+
}
|
| 265 |
+
}
|
| 266 |
+
|
| 267 |
+
inline size_t size_index(size_t size) {
|
| 268 |
+
return c10::llvm::Log2_64_Ceil(size);
|
| 269 |
+
}
|
| 270 |
+
|
| 271 |
+
virtual bool pinned_use_background_threads() {
|
| 272 |
+
return false;
|
| 273 |
+
}
|
| 274 |
+
|
| 275 |
+
virtual void copy_data(void* dest [[maybe_unused]], const void* src [[maybe_unused]], std::size_t count [[maybe_unused]]) const {
|
| 276 |
+
TORCH_CHECK_NOT_IMPLEMENTED(false, "Not implemented for copy_data");
|
| 277 |
+
}
|
| 278 |
+
|
| 279 |
+
private:
|
| 280 |
+
virtual void add_allocated_block(B* block) {
|
| 281 |
+
std::lock_guard<std::mutex> g(blocks_mutex_);
|
| 282 |
+
blocks_.insert(block);
|
| 283 |
+
ptr_to_block_.insert({block->ptr_, block});
|
| 284 |
+
}
|
| 285 |
+
|
| 286 |
+
virtual B* get_free_block(size_t size) {
|
| 287 |
+
auto index = size_index(size);
|
| 288 |
+
std::lock_guard<std::mutex> g(free_list_[index].mutex_);
|
| 289 |
+
if (free_list_[index].list_.size() > 0) {
|
| 290 |
+
B* block = free_list_[index].list_.back();
|
| 291 |
+
free_list_[index].list_.pop_back();
|
| 292 |
+
block->allocated_ = true;
|
| 293 |
+
return block;
|
| 294 |
+
}
|
| 295 |
+
return nullptr;
|
| 296 |
+
}
|
| 297 |
+
|
| 298 |
+
virtual void process_events() {
|
| 299 |
+
// process all events until the last unready event, not for specific size.
|
| 300 |
+
process_events_for_specific_size(-1);
|
| 301 |
+
}
|
| 302 |
+
|
| 303 |
+
// If size is -1, process all events from backwards until the last unready
|
| 304 |
+
// event. Otherwise, process events for a specific size and on first ready block
|
| 305 |
+
// is found, add it to the free list and return.
|
| 306 |
+
virtual void process_events_for_specific_size(int64_t size) {
|
| 307 |
+
size_t event_count = 0;
|
| 308 |
+
size_t max_events = 0;
|
| 309 |
+
{
|
| 310 |
+
std::lock_guard<std::mutex> g(events_mutex_);
|
| 311 |
+
max_events = events_.size();
|
| 312 |
+
}
|
| 313 |
+
|
| 314 |
+
while (true) {
|
| 315 |
+
// Avoid calling cudaEventDestroy while holding a mutex, so move
|
| 316 |
+
// intermediate events out of the lock into this object.
|
| 317 |
+
// process the last event
|
| 318 |
+
std::optional<std::pair<E, B*>> processed;
|
| 319 |
+
{
|
| 320 |
+
std::lock_guard<std::mutex> g(events_mutex_);
|
| 321 |
+
if (!events_.empty()) {
|
| 322 |
+
processed = std::move(events_.back());
|
| 323 |
+
events_.pop_back();
|
| 324 |
+
}
|
| 325 |
+
}
|
| 326 |
+
|
| 327 |
+
if (!processed) {
|
| 328 |
+
return;
|
| 329 |
+
}
|
| 330 |
+
|
| 331 |
+
if (size != -1) {
|
| 332 |
+
if (event_count++ > max_events) {
|
| 333 |
+
{
|
| 334 |
+
std::lock_guard<std::mutex> g(events_mutex_);
|
| 335 |
+
events_.push_front(std::move(*processed));
|
| 336 |
+
}
|
| 337 |
+
return;
|
| 338 |
+
}
|
| 339 |
+
if (size != (int64_t)processed->second->size_) {
|
| 340 |
+
// if we are processing a specific size, and the size of the block
|
| 341 |
+
// doesn't match, we can't use it.
|
| 342 |
+
{
|
| 343 |
+
std::lock_guard<std::mutex> g(events_mutex_);
|
| 344 |
+
events_.push_front(std::move(*processed));
|
| 345 |
+
}
|
| 346 |
+
continue;
|
| 347 |
+
}
|
| 348 |
+
}
|
| 349 |
+
|
| 350 |
+
// otherwise, query the event
|
| 351 |
+
{
|
| 352 |
+
// now, see if we can handle this element
|
| 353 |
+
auto& event = processed->first;
|
| 354 |
+
if (!query_event(event)) {
|
| 355 |
+
// push the event onto the back if it's not ready.
|
| 356 |
+
{
|
| 357 |
+
std::lock_guard<std::mutex> g(events_mutex_);
|
| 358 |
+
if (size == -1) {
|
| 359 |
+
events_.push_back(std::move(*processed));
|
| 360 |
+
return;
|
| 361 |
+
} else {
|
| 362 |
+
events_.push_front(std::move(*processed));
|
| 363 |
+
continue;
|
| 364 |
+
}
|
| 365 |
+
}
|
| 366 |
+
}
|
| 367 |
+
}
|
| 368 |
+
|
| 369 |
+
// Process the events.
|
| 370 |
+
TORCH_INTERNAL_ASSERT(processed);
|
| 371 |
+
auto* block = processed->second;
|
| 372 |
+
bool available = false;
|
| 373 |
+
{
|
| 374 |
+
std::lock_guard<std::mutex> g(block->mutex_);
|
| 375 |
+
TORCH_INTERNAL_ASSERT(!block->allocated_)
|
| 376 |
+
block->event_count_--;
|
| 377 |
+
if (block->event_count_ == 0) {
|
| 378 |
+
available = true;
|
| 379 |
+
}
|
| 380 |
+
}
|
| 381 |
+
|
| 382 |
+
if (available) {
|
| 383 |
+
auto index = size_index(block->size_);
|
| 384 |
+
std::lock_guard<std::mutex> g(free_list_[index].mutex_);
|
| 385 |
+
free_list_[index].list_.push_back(block);
|
| 386 |
+
if (size != -1) {
|
| 387 |
+
return;
|
| 388 |
+
}
|
| 389 |
+
}
|
| 390 |
+
}
|
| 391 |
+
}
|
| 392 |
+
|
| 393 |
+
TaskThreadPool* getBackgroundThreadPool() {
|
| 394 |
+
static TaskThreadPool* pool = new TaskThreadPool(1);
|
| 395 |
+
return pool;
|
| 396 |
+
}
|
| 397 |
+
|
| 398 |
+
/* These following functions are runtime-related. */
|
| 399 |
+
|
| 400 |
+
// Allocate page-locked memory on the host.
|
| 401 |
+
virtual void allocate_host_memory(size_t size, void** ptr) {
|
| 402 |
+
TORCH_CHECK_NOT_IMPLEMENTED(
|
| 403 |
+
false, "Not implemented for allocate_host_memory");
|
| 404 |
+
}
|
| 405 |
+
|
| 406 |
+
// Free block and release the pointer contained in block.
|
| 407 |
+
virtual void free_block(B* block) {
|
| 408 |
+
TORCH_CHECK_NOT_IMPLEMENTED(false, "Not implemented for free_block");
|
| 409 |
+
}
|
| 410 |
+
|
| 411 |
+
// Record an event on stream and store event into events.
|
| 412 |
+
virtual void record_stream(std::optional<std::vector<E>>& events, S stream) {
|
| 413 |
+
TORCH_CHECK_NOT_IMPLEMENTED(false, "Not implemented for record_stream");
|
| 414 |
+
}
|
| 415 |
+
|
| 416 |
+
// Query event if it is completed.
|
| 417 |
+
virtual bool query_event(E& event) {
|
| 418 |
+
TORCH_CHECK_NOT_IMPLEMENTED(false, "Not implemented for query_event");
|
| 419 |
+
}
|
| 420 |
+
|
| 421 |
+
alignas(64) std::mutex blocks_mutex_;
|
| 422 |
+
ska::flat_hash_set<B*> blocks_; // block list
|
| 423 |
+
ska::flat_hash_map<void*, B*> ptr_to_block_;
|
| 424 |
+
|
| 425 |
+
// We keep free list as a vector of free lists, one for each power of two
|
| 426 |
+
// size. This allows us to quickly find a free block of the right size.
|
| 427 |
+
// We use deque to store per size free list and guard the list with its own
|
| 428 |
+
// mutex.
|
| 429 |
+
alignas(64) std::vector<FreeBlockList<B>> free_list_ = std::vector<FreeBlockList<B>>(MAX_SIZE_INDEX);
|
| 430 |
+
|
| 431 |
+
alignas(64) std::mutex events_mutex_;
|
| 432 |
+
std::deque<std::pair<E, B*>> events_; // event queue paired with block
|
| 433 |
+
};
|
| 434 |
+
|
| 435 |
+
template <typename T>
|
| 436 |
+
struct CachingHostAllocatorInterface : public at::Allocator {
|
| 437 |
+
CachingHostAllocatorInterface() : impl_(std::make_unique<T>()) {}
|
| 438 |
+
|
| 439 |
+
at::DataPtr allocate(size_t size) override {
|
| 440 |
+
TORCH_CHECK_NOT_IMPLEMENTED(false, "Not implemented for allocate");
|
| 441 |
+
}
|
| 442 |
+
|
| 443 |
+
void free(void* ctx) {
|
| 444 |
+
impl_->free(ctx);
|
| 445 |
+
}
|
| 446 |
+
|
| 447 |
+
template <typename S>
|
| 448 |
+
bool record_event(void* ptr, void* ctx, S stream) {
|
| 449 |
+
return impl_->record_event(ptr, ctx, stream);
|
| 450 |
+
}
|
| 451 |
+
|
| 452 |
+
void empty_cache() {
|
| 453 |
+
impl_->empty_cache();
|
| 454 |
+
}
|
| 455 |
+
|
| 456 |
+
void copy_data(void* dest, const void* src, std::size_t count)
|
| 457 |
+
const override {
|
| 458 |
+
impl_->copy_data(dest, src, count);
|
| 459 |
+
}
|
| 460 |
+
|
| 461 |
+
std::unique_ptr<T> impl_;
|
| 462 |
+
};
|
| 463 |
+
|
| 464 |
+
} // namespace at
|
| 465 |
+
C10_DIAGNOSTIC_POP()
|
lib/python3.12/site-packages/torch/include/ATen/core/CheckMemoryFormat.h
ADDED
|
@@ -0,0 +1,24 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#include <c10/core/TensorOptions.h>
|
| 2 |
+
|
| 3 |
+
namespace c10::impl {
|
| 4 |
+
|
| 5 |
+
inline std::optional<MemoryFormat>
|
| 6 |
+
check_tensor_options_and_extract_memory_format(
|
| 7 |
+
const TensorOptions& options,
|
| 8 |
+
std::optional<MemoryFormat> memory_format) {
|
| 9 |
+
TORCH_CHECK(
|
| 10 |
+
options.requires_grad_opt() != true,
|
| 11 |
+
"Operators taking TensorOptions cannot take a TensorOptions with "
|
| 12 |
+
"options.requires_grad set as true. This isn't implemented yet.");
|
| 13 |
+
TORCH_CHECK(
|
| 14 |
+
!(options.has_memory_format() && memory_format.has_value()),
|
| 15 |
+
"Cannot set memory_format both in TensorOptions and explicit argument; please delete "
|
| 16 |
+
"the redundant setter.");
|
| 17 |
+
if (memory_format.has_value()) {
|
| 18 |
+
return memory_format;
|
| 19 |
+
} else {
|
| 20 |
+
return options.memory_format_opt();
|
| 21 |
+
}
|
| 22 |
+
}
|
| 23 |
+
|
| 24 |
+
} // namespace impl namespace c10
|
lib/python3.12/site-packages/torch/include/ATen/core/DeprecatedTypeProperties.h
ADDED
|
@@ -0,0 +1,139 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <c10/core/Backend.h>
|
| 4 |
+
#include <c10/core/ScalarType.h>
|
| 5 |
+
#include <c10/core/Layout.h>
|
| 6 |
+
#include <c10/core/TensorOptions.h>
|
| 7 |
+
#include <c10/core/Storage.h>
|
| 8 |
+
#include <ATen/core/DeprecatedTypePropertiesRegistry.h>
|
| 9 |
+
#include <ATen/core/Generator.h>
|
| 10 |
+
|
| 11 |
+
|
| 12 |
+
namespace at {
|
| 13 |
+
|
| 14 |
+
class Tensor;
|
| 15 |
+
|
| 16 |
+
// This class specifies a Backend and a ScalarType. Currently, it primarily
|
| 17 |
+
// serves as a replacement return value for Tensor::type(). Previously,
|
| 18 |
+
// Tensor::type() returned Type&, but we are changing Type to not be
|
| 19 |
+
// dtype-specific.
|
| 20 |
+
class TORCH_API DeprecatedTypeProperties {
|
| 21 |
+
public:
|
| 22 |
+
DeprecatedTypeProperties(Backend backend, ScalarType scalar_type)
|
| 23 |
+
: backend_(backend), scalar_type_(scalar_type) {}
|
| 24 |
+
|
| 25 |
+
Backend backend() const {
|
| 26 |
+
return backend_;
|
| 27 |
+
}
|
| 28 |
+
|
| 29 |
+
Layout layout() const {
|
| 30 |
+
return layout_from_backend(backend_);
|
| 31 |
+
}
|
| 32 |
+
|
| 33 |
+
bool is_sparse() const {
|
| 34 |
+
return layout_from_backend(backend()) == kSparse;
|
| 35 |
+
}
|
| 36 |
+
|
| 37 |
+
bool is_sparse_csr() const {
|
| 38 |
+
return layout_from_backend(backend()) == kSparseCsr;
|
| 39 |
+
}
|
| 40 |
+
|
| 41 |
+
c10::DeviceType device_type() const {
|
| 42 |
+
return backendToDeviceType(backend_);
|
| 43 |
+
}
|
| 44 |
+
|
| 45 |
+
bool is_cuda() const {
|
| 46 |
+
return backendToDeviceType(backend_) == kCUDA;
|
| 47 |
+
}
|
| 48 |
+
|
| 49 |
+
ScalarType scalarType() const {
|
| 50 |
+
return scalar_type_;
|
| 51 |
+
}
|
| 52 |
+
|
| 53 |
+
caffe2::TypeMeta typeMeta() const {
|
| 54 |
+
return scalarTypeToTypeMeta(scalar_type_);
|
| 55 |
+
}
|
| 56 |
+
|
| 57 |
+
bool operator==(const DeprecatedTypeProperties& other) const {
|
| 58 |
+
return backend_ == other.backend() && scalar_type_ == other.scalarType();
|
| 59 |
+
}
|
| 60 |
+
|
| 61 |
+
bool operator!=(const DeprecatedTypeProperties& other) const {
|
| 62 |
+
return !(*this == other);
|
| 63 |
+
}
|
| 64 |
+
|
| 65 |
+
std::string toString() const {
|
| 66 |
+
std::string base_str;
|
| 67 |
+
if (backend_ == Backend::Undefined || scalar_type_ == ScalarType::Undefined) {
|
| 68 |
+
base_str = "UndefinedType";
|
| 69 |
+
} else {
|
| 70 |
+
base_str = std::string(at::toString(backend_)) + at::toString(scalar_type_) + "Type";
|
| 71 |
+
}
|
| 72 |
+
return base_str;
|
| 73 |
+
}
|
| 74 |
+
|
| 75 |
+
DeprecatedTypeProperties & toBackend(Backend b) const {
|
| 76 |
+
return globalDeprecatedTypePropertiesRegistry().getDeprecatedTypeProperties(
|
| 77 |
+
b, scalar_type_);
|
| 78 |
+
}
|
| 79 |
+
|
| 80 |
+
DeprecatedTypeProperties & toScalarType(ScalarType s) const {
|
| 81 |
+
return globalDeprecatedTypePropertiesRegistry().getDeprecatedTypeProperties(
|
| 82 |
+
backend_, s);
|
| 83 |
+
}
|
| 84 |
+
|
| 85 |
+
DeprecatedTypeProperties & cpu() const {
|
| 86 |
+
return toBackend(Backend::CPU);
|
| 87 |
+
}
|
| 88 |
+
|
| 89 |
+
DeprecatedTypeProperties & cuda() const {
|
| 90 |
+
return toBackend(Backend::CUDA);
|
| 91 |
+
}
|
| 92 |
+
|
| 93 |
+
DeprecatedTypeProperties & hip() const {
|
| 94 |
+
return toBackend(Backend::HIP);
|
| 95 |
+
}
|
| 96 |
+
|
| 97 |
+
DeprecatedTypeProperties & privateUser1() const {
|
| 98 |
+
return toBackend(Backend::PrivateUse1);
|
| 99 |
+
}
|
| 100 |
+
|
| 101 |
+
/// Constructs the `TensorOptions` from a type and a `device_index`.
|
| 102 |
+
TensorOptions options(int16_t device_index = -1) const {
|
| 103 |
+
return TensorOptions().dtype(typeMeta())
|
| 104 |
+
.device(device_type(), static_cast<c10::DeviceIndex>(device_index))
|
| 105 |
+
.layout(layout());
|
| 106 |
+
}
|
| 107 |
+
|
| 108 |
+
/// Constructs the `TensorOptions` from a type and a Device. Asserts that
|
| 109 |
+
/// the device type matches the device type of the type.
|
| 110 |
+
TensorOptions options(std::optional<Device> device_opt) const {
|
| 111 |
+
if (!device_opt.has_value()) {
|
| 112 |
+
return options(-1);
|
| 113 |
+
} else {
|
| 114 |
+
Device device = device_opt.value();
|
| 115 |
+
AT_ASSERT(device.type() == device_type());
|
| 116 |
+
return options(device.index());
|
| 117 |
+
}
|
| 118 |
+
}
|
| 119 |
+
|
| 120 |
+
operator TensorOptions() const {
|
| 121 |
+
return options();
|
| 122 |
+
}
|
| 123 |
+
|
| 124 |
+
int64_t id() const {
|
| 125 |
+
return static_cast<int64_t>(backend()) *
|
| 126 |
+
static_cast<int64_t>(ScalarType::NumOptions) +
|
| 127 |
+
static_cast<int64_t>(scalarType());
|
| 128 |
+
}
|
| 129 |
+
|
| 130 |
+
Tensor unsafeTensorFromTH(void * th_pointer, bool retain) const;
|
| 131 |
+
Storage unsafeStorageFromTH(void * th_pointer, bool retain) const;
|
| 132 |
+
Tensor copy(const Tensor & src, bool non_blocking=false, std::optional<Device> to_device={}) const;
|
| 133 |
+
|
| 134 |
+
private:
|
| 135 |
+
Backend backend_;
|
| 136 |
+
ScalarType scalar_type_;
|
| 137 |
+
};
|
| 138 |
+
|
| 139 |
+
} // namespace at
|
lib/python3.12/site-packages/torch/include/ATen/core/DeprecatedTypePropertiesRegistry.h
ADDED
|
@@ -0,0 +1,33 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
// In order to preserve bc, we make DeprecatedTypeProperties instances unique
|
| 4 |
+
// just like they are for Type.
|
| 5 |
+
|
| 6 |
+
#include <c10/core/Backend.h>
|
| 7 |
+
#include <c10/core/ScalarType.h>
|
| 8 |
+
#include <memory>
|
| 9 |
+
|
| 10 |
+
namespace at {
|
| 11 |
+
|
| 12 |
+
class DeprecatedTypeProperties;
|
| 13 |
+
|
| 14 |
+
struct TORCH_API DeprecatedTypePropertiesDeleter {
|
| 15 |
+
void operator()(DeprecatedTypeProperties * ptr);
|
| 16 |
+
};
|
| 17 |
+
|
| 18 |
+
class TORCH_API DeprecatedTypePropertiesRegistry {
|
| 19 |
+
public:
|
| 20 |
+
DeprecatedTypePropertiesRegistry();
|
| 21 |
+
|
| 22 |
+
DeprecatedTypeProperties& getDeprecatedTypeProperties(Backend p, ScalarType s) const;
|
| 23 |
+
|
| 24 |
+
private:
|
| 25 |
+
// NOLINTNEXTLINE(*c-array*)
|
| 26 |
+
std::unique_ptr<DeprecatedTypeProperties> registry
|
| 27 |
+
[static_cast<int>(Backend::NumOptions)]
|
| 28 |
+
[static_cast<int>(ScalarType::NumOptions)];
|
| 29 |
+
};
|
| 30 |
+
|
| 31 |
+
TORCH_API DeprecatedTypePropertiesRegistry& globalDeprecatedTypePropertiesRegistry();
|
| 32 |
+
|
| 33 |
+
} // namespace at
|
lib/python3.12/site-packages/torch/include/ATen/core/Dict.h
ADDED
|
@@ -0,0 +1,399 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <c10/macros/Macros.h>
|
| 4 |
+
#include <c10/macros/Export.h>
|
| 5 |
+
#include <c10/util/TypeTraits.h>
|
| 6 |
+
#include <c10/util/TypeList.h>
|
| 7 |
+
#include <c10/util/intrusive_ptr.h>
|
| 8 |
+
#include <c10/util/order_preserving_flat_hash_map.h>
|
| 9 |
+
#include <optional>
|
| 10 |
+
#include <ATen/core/TensorBody.h>
|
| 11 |
+
#include <ATen/core/jit_type_base.h>
|
| 12 |
+
|
| 13 |
+
namespace c10 {
|
| 14 |
+
struct IValue;
|
| 15 |
+
template<class Key, class Value> class Dict;
|
| 16 |
+
struct Type;
|
| 17 |
+
|
| 18 |
+
namespace impl {
|
| 19 |
+
|
| 20 |
+
using valid_dict_key_types = guts::typelist::typelist<
|
| 21 |
+
int64_t,
|
| 22 |
+
std::string,
|
| 23 |
+
double,
|
| 24 |
+
c10::complex<double>,
|
| 25 |
+
bool,
|
| 26 |
+
at::Tensor
|
| 27 |
+
>;
|
| 28 |
+
}
|
| 29 |
+
|
| 30 |
+
namespace detail {
|
| 31 |
+
|
| 32 |
+
struct DictKeyHash {
|
| 33 |
+
size_t operator()(const IValue& ivalue) const;
|
| 34 |
+
};
|
| 35 |
+
|
| 36 |
+
struct DictKeyEqualTo {
|
| 37 |
+
bool operator()(const IValue& lhs, const IValue& rhs) const;
|
| 38 |
+
};
|
| 39 |
+
|
| 40 |
+
struct DictImpl final : public c10::intrusive_ptr_target {
|
| 41 |
+
using dict_map_type = ska_ordered::order_preserving_flat_hash_map<IValue, IValue, DictKeyHash, DictKeyEqualTo>;
|
| 42 |
+
struct DictElementTypes final {
|
| 43 |
+
TypePtr keyType;
|
| 44 |
+
TypePtr valueType;
|
| 45 |
+
};
|
| 46 |
+
|
| 47 |
+
explicit DictImpl(dict_map_type dict_, DictElementTypes elementTypes_)
|
| 48 |
+
: dict(std::move(dict_))
|
| 49 |
+
, elementTypes(std::move(elementTypes_)) {}
|
| 50 |
+
dict_map_type dict;
|
| 51 |
+
|
| 52 |
+
DictElementTypes elementTypes;
|
| 53 |
+
|
| 54 |
+
intrusive_ptr<DictImpl> copy() const;
|
| 55 |
+
friend TORCH_API bool operator==(const DictImpl& lhs, const DictImpl& rhs);
|
| 56 |
+
};
|
| 57 |
+
|
| 58 |
+
}
|
| 59 |
+
|
| 60 |
+
namespace impl {
|
| 61 |
+
template<class Key, class Value, class Iterator> class DictIterator;
|
| 62 |
+
|
| 63 |
+
/**
|
| 64 |
+
* A reference to an entry in the Dict.
|
| 65 |
+
* Use the `key()` and `value()` methods to read the element.
|
| 66 |
+
*/
|
| 67 |
+
template<class Key, class Value, class Iterator>
|
| 68 |
+
class DictEntryRef final {
|
| 69 |
+
public:
|
| 70 |
+
explicit DictEntryRef(Iterator iterator)
|
| 71 |
+
: iterator_(std::move(iterator)) {}
|
| 72 |
+
|
| 73 |
+
decltype(auto) key() const {
|
| 74 |
+
return iterator_->first.template to<Key>();
|
| 75 |
+
}
|
| 76 |
+
|
| 77 |
+
decltype(auto) value() const {
|
| 78 |
+
return iterator_->second.template to<Value>();
|
| 79 |
+
}
|
| 80 |
+
|
| 81 |
+
template<class Value_>
|
| 82 |
+
void setValue(Value_&& value) const {
|
| 83 |
+
static_assert(std::is_constructible_v<Value, Value_>, "Wrong type for the value argument of setValue()");
|
| 84 |
+
iterator_->second = Value(std::forward<Value_>(value));
|
| 85 |
+
}
|
| 86 |
+
~DictEntryRef() = default;
|
| 87 |
+
|
| 88 |
+
private:
|
| 89 |
+
// allow copying and moving, but only our friends (i.e. the Dict class) can do
|
| 90 |
+
// it. Copying/moving this reference wrapper would be too ambiguous to allow it
|
| 91 |
+
// in the public API.
|
| 92 |
+
DictEntryRef(const DictEntryRef&) = default;
|
| 93 |
+
DictEntryRef& operator=(const DictEntryRef&) = default;
|
| 94 |
+
DictEntryRef(DictEntryRef&&) noexcept = default;
|
| 95 |
+
DictEntryRef& operator=(DictEntryRef&& rhs) & noexcept = default;
|
| 96 |
+
|
| 97 |
+
Iterator iterator_;
|
| 98 |
+
friend class DictIterator<Key, Value, Iterator>;
|
| 99 |
+
friend class Dict<Key, Value>;
|
| 100 |
+
};
|
| 101 |
+
|
| 102 |
+
// this wraps map_type::iterator to make sure user code can't rely
|
| 103 |
+
// on it being the type of the underlying map.
|
| 104 |
+
template<class Key, class Value, class Iterator>
|
| 105 |
+
class DictIterator final {
|
| 106 |
+
public:
|
| 107 |
+
// C++17 friendly std::iterator implementation
|
| 108 |
+
using iterator_category = std::forward_iterator_tag;
|
| 109 |
+
using value_type = DictEntryRef<Key, Value, Iterator>;
|
| 110 |
+
using difference_type = std::ptrdiff_t;
|
| 111 |
+
using pointer = value_type*;
|
| 112 |
+
using reference = value_type&;
|
| 113 |
+
|
| 114 |
+
explicit DictIterator() = default;
|
| 115 |
+
~DictIterator() = default;
|
| 116 |
+
|
| 117 |
+
DictIterator(const DictIterator& rhs): entryRef_(rhs.entryRef_) {}
|
| 118 |
+
DictIterator(DictIterator&& rhs) noexcept: entryRef_(std::move(rhs.entryRef_)) {}
|
| 119 |
+
DictIterator& operator=(const DictIterator& rhs) {
|
| 120 |
+
entryRef_ = rhs.entryRef_;
|
| 121 |
+
return *this;
|
| 122 |
+
}
|
| 123 |
+
DictIterator& operator=(DictIterator&& rhs) noexcept {
|
| 124 |
+
entryRef_ = std::move(rhs.entryRef_);
|
| 125 |
+
return *this;
|
| 126 |
+
}
|
| 127 |
+
|
| 128 |
+
DictIterator& operator++() {
|
| 129 |
+
++entryRef_.iterator_;
|
| 130 |
+
return *this;
|
| 131 |
+
}
|
| 132 |
+
|
| 133 |
+
DictIterator operator++(int) {
|
| 134 |
+
DictIterator copy(*this);
|
| 135 |
+
++*this;
|
| 136 |
+
return copy;
|
| 137 |
+
}
|
| 138 |
+
|
| 139 |
+
const DictEntryRef<Key, Value, Iterator>& operator*() const {
|
| 140 |
+
return entryRef_;
|
| 141 |
+
}
|
| 142 |
+
|
| 143 |
+
const DictEntryRef<Key, Value, Iterator>* operator->() const {
|
| 144 |
+
return &entryRef_;
|
| 145 |
+
}
|
| 146 |
+
|
| 147 |
+
friend difference_type operator-(const DictIterator& lhs, const DictIterator& rhs) {
|
| 148 |
+
return lhs.entryRef_.iterator_ - rhs.entryRef_.iterator_;
|
| 149 |
+
}
|
| 150 |
+
|
| 151 |
+
private:
|
| 152 |
+
explicit DictIterator(Iterator iterator): entryRef_(std::move(iterator)) {}
|
| 153 |
+
|
| 154 |
+
const Iterator& get_iterator_() const {
|
| 155 |
+
return entryRef_.iterator_;
|
| 156 |
+
}
|
| 157 |
+
|
| 158 |
+
friend bool operator==(const DictIterator& lhs, const DictIterator& rhs) {
|
| 159 |
+
return lhs.get_iterator_() == rhs.get_iterator_();
|
| 160 |
+
}
|
| 161 |
+
|
| 162 |
+
friend bool operator!=(const DictIterator& lhs, const DictIterator& rhs) {
|
| 163 |
+
return lhs.get_iterator_() != rhs.get_iterator_();
|
| 164 |
+
}
|
| 165 |
+
|
| 166 |
+
friend bool operator<(const DictIterator& lhs, const DictIterator& rhs) {
|
| 167 |
+
return lhs.get_iterator_() < rhs.get_iterator_();
|
| 168 |
+
}
|
| 169 |
+
|
| 170 |
+
friend bool operator<=(const DictIterator& lhs, const DictIterator& rhs) {
|
| 171 |
+
return lhs.get_iterator_() <= rhs.get_iterator_();
|
| 172 |
+
}
|
| 173 |
+
|
| 174 |
+
friend bool operator>(const DictIterator& lhs, const DictIterator& rhs) {
|
| 175 |
+
return lhs.get_iterator_() > rhs.get_iterator_();
|
| 176 |
+
}
|
| 177 |
+
|
| 178 |
+
friend bool operator>=(const DictIterator& lhs, const DictIterator& rhs) {
|
| 179 |
+
return lhs.get_iterator_() >= rhs.get_iterator_();
|
| 180 |
+
}
|
| 181 |
+
|
| 182 |
+
DictEntryRef<Key, Value, Iterator> entryRef_;
|
| 183 |
+
|
| 184 |
+
friend class DictIterator<Key, Value, typename c10::detail::DictImpl::dict_map_type::iterator>;
|
| 185 |
+
friend class Dict<Key, Value>;
|
| 186 |
+
};
|
| 187 |
+
|
| 188 |
+
template<class Key, class Value> Dict<Key, Value> toTypedDict(Dict<IValue, IValue> dict);
|
| 189 |
+
template<class Key, class Value> Dict<IValue, IValue> toGenericDict(Dict<Key, Value> dict);
|
| 190 |
+
}
|
| 191 |
+
|
| 192 |
+
/**
|
| 193 |
+
* An object of this class stores a map from Key to Value.
|
| 194 |
+
*
|
| 195 |
+
* This is a pointer type. After a copy, both Dicts
|
| 196 |
+
* will share the same storage:
|
| 197 |
+
*
|
| 198 |
+
* > Dict<int, string> a;
|
| 199 |
+
* > Dict<int, string> b = a;
|
| 200 |
+
* > b.insert(3, "three");
|
| 201 |
+
* > ASSERT("three" == a.at(3));
|
| 202 |
+
*
|
| 203 |
+
* We use this class in the PyTorch kernel API because that
|
| 204 |
+
* allows us to do optimizations and switch out the underlying
|
| 205 |
+
* map implementation without breaking backwards compatibility
|
| 206 |
+
* for the kernel API.
|
| 207 |
+
*/
|
| 208 |
+
template<class Key, class Value>
|
| 209 |
+
// NOLINTNEXTLINE(cppcoreguidelines-special-member-functions)
|
| 210 |
+
class Dict final {
|
| 211 |
+
private:
|
| 212 |
+
static_assert((std::is_same_v<IValue, Key> && std::is_same_v<IValue, Value>) || guts::typelist::contains<impl::valid_dict_key_types, Key>::value, "Invalid Key type for Dict. We only support int64_t, double, bool, and string.");
|
| 213 |
+
|
| 214 |
+
// impl_ stores the underlying map as a ska_ordered::order_preserving_flat_hash_map.
|
| 215 |
+
// We intentionally don't offer conversion from/to
|
| 216 |
+
// order_preserving_flat_hash_map, return references to it or something like that,
|
| 217 |
+
// because such operations would get expensive if we switch out
|
| 218 |
+
// the actual map implementation.
|
| 219 |
+
// This is an intrusive_ptr because Dict is a pointer type.
|
| 220 |
+
// Invariant: This will never be a nullptr, there will always be a valid
|
| 221 |
+
// DictImpl.
|
| 222 |
+
c10::intrusive_ptr<detail::DictImpl> impl_;
|
| 223 |
+
|
| 224 |
+
explicit Dict(c10::intrusive_ptr<detail::DictImpl>&& impl);
|
| 225 |
+
friend struct IValue;
|
| 226 |
+
template<class K, class V> friend Dict<K, V> impl::toTypedDict(Dict<IValue, IValue>);
|
| 227 |
+
template<class K, class V> friend Dict<IValue, IValue> impl::toGenericDict(Dict<K, V>);
|
| 228 |
+
|
| 229 |
+
public:
|
| 230 |
+
using key_type = Key;
|
| 231 |
+
using mapped_type = Value;
|
| 232 |
+
using size_type = typename detail::DictImpl::dict_map_type::size_type;
|
| 233 |
+
using iterator = impl::DictIterator<Key, Value, typename detail::DictImpl::dict_map_type::iterator>;
|
| 234 |
+
|
| 235 |
+
/**
|
| 236 |
+
* Creates an empty dict.
|
| 237 |
+
*/
|
| 238 |
+
explicit Dict();
|
| 239 |
+
|
| 240 |
+
/**
|
| 241 |
+
* Create a generic dict with runtime type information.
|
| 242 |
+
* This only works for c10::impl::GenericDict and is not part of the public API
|
| 243 |
+
* but only supposed to be used internally by PyTorch.
|
| 244 |
+
*/
|
| 245 |
+
explicit Dict(TypePtr keyType, TypePtr valueType);
|
| 246 |
+
|
| 247 |
+
~Dict() = default;
|
| 248 |
+
|
| 249 |
+
Dict(const Dict&) = default;
|
| 250 |
+
Dict& operator=(const Dict&) = default;
|
| 251 |
+
|
| 252 |
+
/**
|
| 253 |
+
* Create a new Dict pointing to a deep copy of the same data.
|
| 254 |
+
* The Dict returned is a new dict with separate storage.
|
| 255 |
+
* Changes in it are not reflected in the original dict or vice versa.
|
| 256 |
+
*/
|
| 257 |
+
Dict copy() const;
|
| 258 |
+
|
| 259 |
+
/**
|
| 260 |
+
* Returns an iterator to the first element of the container.
|
| 261 |
+
* If the container is empty, the returned iterator will be equal to end().
|
| 262 |
+
*/
|
| 263 |
+
iterator begin() const;
|
| 264 |
+
|
| 265 |
+
/**
|
| 266 |
+
* Returns an iterator to the element following the last element of the container.
|
| 267 |
+
* This element acts as a placeholder; attempting to access it results in undefined behavior.
|
| 268 |
+
*/
|
| 269 |
+
iterator end() const;
|
| 270 |
+
|
| 271 |
+
/**
|
| 272 |
+
* Checks if the container has no elements.
|
| 273 |
+
*/
|
| 274 |
+
bool empty() const;
|
| 275 |
+
|
| 276 |
+
/**
|
| 277 |
+
* Returns the number of elements in the container.
|
| 278 |
+
*/
|
| 279 |
+
size_type size() const;
|
| 280 |
+
|
| 281 |
+
/**
|
| 282 |
+
* Erases all elements from the container. After this call, size() returns zero.
|
| 283 |
+
* Invalidates any references, pointers, or iterators referring to contained elements. May also invalidate past-the-end iterators.
|
| 284 |
+
*/
|
| 285 |
+
void clear() const;
|
| 286 |
+
|
| 287 |
+
/**
|
| 288 |
+
* Inserts element(s) into the container, if the container doesn't already contain an element with an equivalent key.
|
| 289 |
+
* May invalidate any references, pointers, or iterators referring to contained elements.
|
| 290 |
+
*
|
| 291 |
+
* @return A pair consisting of an iterator to the inserted element (or to the element that prevented the insertion) and a bool denoting whether the insertion took place.
|
| 292 |
+
*/
|
| 293 |
+
template<class Key_, class Value_>
|
| 294 |
+
std::pair<iterator, bool> insert(Key_&& key, Value_&& value) const;
|
| 295 |
+
|
| 296 |
+
/**
|
| 297 |
+
* If an element with the given key already exists, it is overwritten with the given value.
|
| 298 |
+
* Otherwise, a new element with the given key and value are inserted.
|
| 299 |
+
* May invalidate any references, pointers, or iterators referring to contained elements.
|
| 300 |
+
*
|
| 301 |
+
* @return The bool component is true if the insertion took place and false if the assignment took place. The iterator component is pointing at the element that was inserted or updated.
|
| 302 |
+
*/
|
| 303 |
+
template<class Key_, class Value_>
|
| 304 |
+
std::pair<iterator, bool> insert_or_assign(Key_&& key, Value_&& value) const;
|
| 305 |
+
|
| 306 |
+
/**
|
| 307 |
+
* Removes the element pointed to by iter.
|
| 308 |
+
* May invalidate any references, pointers, or iterators referring to contained elements.
|
| 309 |
+
* The iterator iter must be valid and dereferenceable. Thus the end() iterator (which is valid, but is not dereferenceable) cannot be used as a value for iter.
|
| 310 |
+
*/
|
| 311 |
+
void erase(iterator iter) const;
|
| 312 |
+
|
| 313 |
+
/**
|
| 314 |
+
* Removes the element with the given key, if it exists.
|
| 315 |
+
* May invalidate any references, pointers, or iterators referring to contained elements.
|
| 316 |
+
*
|
| 317 |
+
* @return The number of elements removed. This is either '1' if an element with the key existed, or '0' if it didn't.
|
| 318 |
+
*/
|
| 319 |
+
[[nodiscard]] size_t erase(const Key& key) const;
|
| 320 |
+
|
| 321 |
+
/**
|
| 322 |
+
* Returns the mapped value of the element with key equivalent to key.
|
| 323 |
+
* If no such element exists, an exception of type std::out_of_range is thrown.
|
| 324 |
+
*/
|
| 325 |
+
Value at(const Key& key) const;
|
| 326 |
+
|
| 327 |
+
/**
|
| 328 |
+
* Finds an element with key equivalent to key.
|
| 329 |
+
*
|
| 330 |
+
* @return Iterator to an element with key equivalent to key.
|
| 331 |
+
* If no such element is found, past-the-end (see end()) iterator is returned.
|
| 332 |
+
*/
|
| 333 |
+
iterator find(const Key& key) const;
|
| 334 |
+
|
| 335 |
+
/**
|
| 336 |
+
* Checks if there is an element with key equivalent to key in the container.
|
| 337 |
+
*
|
| 338 |
+
* @return true if there is such an element, otherwise false.
|
| 339 |
+
*/
|
| 340 |
+
bool contains(const Key& key) const;
|
| 341 |
+
|
| 342 |
+
/**
|
| 343 |
+
* Increase the capacity so that at least count elements can be stored without
|
| 344 |
+
* having to reallocate or rehash.
|
| 345 |
+
*/
|
| 346 |
+
void reserve(size_type count) const;
|
| 347 |
+
|
| 348 |
+
/**
|
| 349 |
+
* Value equality comparison. This function implements Python-like semantics for
|
| 350 |
+
* equality: two dicts with the same identity (e.g. same pointer) trivially
|
| 351 |
+
* compare equal, otherwise each element is compared for equality.
|
| 352 |
+
*/
|
| 353 |
+
template <class Key_, class Value_>
|
| 354 |
+
friend bool operator==(
|
| 355 |
+
const Dict<Key_, Value_>& lhs,
|
| 356 |
+
const Dict<Key_, Value_>& rhs);
|
| 357 |
+
template <class Key_, class Value_>
|
| 358 |
+
friend bool operator!=(
|
| 359 |
+
const Dict<Key_, Value_>& lhs,
|
| 360 |
+
const Dict<Key_, Value_>& rhs);
|
| 361 |
+
|
| 362 |
+
/**
|
| 363 |
+
* Identity comparison. Returns true if and only if `rhs` represents the same
|
| 364 |
+
* Dict object as `this`.
|
| 365 |
+
*/
|
| 366 |
+
bool is(const Dict& rhs) const;
|
| 367 |
+
|
| 368 |
+
// private API for now because the return type will change to TypePtr
|
| 369 |
+
// instead of std::optional<TypePtr> once types are mandatory.
|
| 370 |
+
TypePtr keyType() const;
|
| 371 |
+
TypePtr valueType() const;
|
| 372 |
+
|
| 373 |
+
// [unsafe set type]
|
| 374 |
+
// These functions mutate the tagged type of this dictionary in place.
|
| 375 |
+
// There is no checking that the members of the dictionary are instances
|
| 376 |
+
// of the new types, nor is there a check that other IValues which
|
| 377 |
+
// hold references to this dictionary have the right static type.
|
| 378 |
+
// This functionality is used only in the unpickler, where at
|
| 379 |
+
// creation type the real type of the dictionary is unknown, but
|
| 380 |
+
// then later recovered from the static type information of the
|
| 381 |
+
// unpickled object.
|
| 382 |
+
void unsafeSetKeyType(TypePtr t);
|
| 383 |
+
void unsafeSetValueType(TypePtr t);
|
| 384 |
+
};
|
| 385 |
+
|
| 386 |
+
namespace impl {
|
| 387 |
+
// GenericDict is how IValue stores dicts. It is, however, not part of the
|
| 388 |
+
// public API. Kernels should use Dicts with concrete Key, Value types instead
|
| 389 |
+
// (maybe except for some internal prim ops).
|
| 390 |
+
using GenericDict = Dict<IValue, IValue>;
|
| 391 |
+
|
| 392 |
+
}
|
| 393 |
+
}
|
| 394 |
+
|
| 395 |
+
namespace torch {
|
| 396 |
+
template<class Key, class Value> using Dict = c10::Dict<Key, Value>;
|
| 397 |
+
}
|
| 398 |
+
|
| 399 |
+
#include <ATen/core/Dict_inl.h> // IWYU pragma: keep
|
lib/python3.12/site-packages/torch/include/ATen/core/Dict_inl.h
ADDED
|
@@ -0,0 +1,209 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <ATen/core/ivalue.h>
|
| 4 |
+
#include <c10/util/hash.h>
|
| 5 |
+
|
| 6 |
+
namespace c10 {
|
| 7 |
+
namespace detail {
|
| 8 |
+
inline bool DictKeyEqualTo::operator()(const IValue& lhs, const IValue& rhs) const {
|
| 9 |
+
if (lhs.isTensor() && rhs.isTensor()) {
|
| 10 |
+
// for tensors, we compare only by identity (following how it's done in Python).
|
| 11 |
+
return lhs.is(rhs);
|
| 12 |
+
}
|
| 13 |
+
// Otherwise, we first compare by identity for efficiency, then by value (see:
|
| 14 |
+
// [container equality])
|
| 15 |
+
return _fastEqualsForContainer(lhs, rhs);
|
| 16 |
+
}
|
| 17 |
+
}
|
| 18 |
+
|
| 19 |
+
template<class T> decltype(auto) getTypePtr();
|
| 20 |
+
std::string toString(const Type& type);
|
| 21 |
+
|
| 22 |
+
namespace impl {
|
| 23 |
+
|
| 24 |
+
template<class Key, class Value>
|
| 25 |
+
Dict<Key, Value> toTypedDict(GenericDict dict) {
|
| 26 |
+
TORCH_INTERNAL_ASSERT(*getTypePtr<Key>() == *dict.impl_->elementTypes.keyType, "Tried to cast a Dict<", toString(*dict.impl_->elementTypes.keyType), ", ", toString(*dict.impl_->elementTypes.valueType) ,"> to a Dict<", toString(*getTypePtr<Key>()), ", ", toString(*getTypePtr<Value>()), ">. Key types mismatch.");
|
| 27 |
+
TORCH_INTERNAL_ASSERT(*getTypePtr<Value>() == *dict.impl_->elementTypes.valueType, "Tried to cast a Dict<", toString(*dict.impl_->elementTypes.keyType), ", ", toString(*dict.impl_->elementTypes.valueType) ,"> to a Dict<", toString(*getTypePtr<Key>()), ", ", toString(*getTypePtr<Value>()), ">. Value types mismatch.");
|
| 28 |
+
|
| 29 |
+
return Dict<Key, Value>(std::move(dict.impl_));
|
| 30 |
+
}
|
| 31 |
+
|
| 32 |
+
template<class Key, class Value>
|
| 33 |
+
GenericDict toGenericDict(Dict<Key, Value> dict) {
|
| 34 |
+
return GenericDict(std::move(dict.impl_));
|
| 35 |
+
}
|
| 36 |
+
}
|
| 37 |
+
|
| 38 |
+
namespace detail {
|
| 39 |
+
|
| 40 |
+
inline size_t DictKeyHash::operator()(const IValue& ivalue) const {
|
| 41 |
+
if (ivalue.isInt()) {
|
| 42 |
+
return std::hash<int64_t>()(ivalue.toInt());
|
| 43 |
+
} else if (ivalue.isString()) {
|
| 44 |
+
return std::hash<std::string_view>()(ivalue.toStringView());
|
| 45 |
+
} else if (ivalue.isDouble()) {
|
| 46 |
+
return std::hash<double>()(ivalue.toDouble());
|
| 47 |
+
} else if (ivalue.isComplexDouble()) {
|
| 48 |
+
return c10::hash<c10::complex<double>>()(ivalue.toComplexDouble());
|
| 49 |
+
} else if (ivalue.isBool()) {
|
| 50 |
+
return std::hash<bool>()(ivalue.toBool());
|
| 51 |
+
} else if (ivalue.isTensor()) {
|
| 52 |
+
return std::hash<TensorImpl*>()(ivalue.toTensor().unsafeGetTensorImpl());
|
| 53 |
+
} else if (ivalue.isDevice()) {
|
| 54 |
+
return std::hash<Device>()(ivalue.toDevice());
|
| 55 |
+
} else {
|
| 56 |
+
throw std::runtime_error(
|
| 57 |
+
"Can't hash IValues with tag '" + ivalue.tagKind() + "'");
|
| 58 |
+
}
|
| 59 |
+
}
|
| 60 |
+
|
| 61 |
+
inline intrusive_ptr<DictImpl> DictImpl::copy() const {
|
| 62 |
+
return make_intrusive<DictImpl>(dict, elementTypes);
|
| 63 |
+
}
|
| 64 |
+
|
| 65 |
+
}
|
| 66 |
+
|
| 67 |
+
template<class Key, class Value>
|
| 68 |
+
Dict<Key, Value>::Dict()
|
| 69 |
+
:Dict(make_intrusive<detail::DictImpl>(
|
| 70 |
+
detail::DictImpl::dict_map_type(),
|
| 71 |
+
detail::DictImpl::DictElementTypes{getTypePtr<Key>(), getTypePtr<Value>()})) {
|
| 72 |
+
static_assert(!std::is_same_v<Key, IValue>, "This constructor is not valid for Dict<IValue, _>. Please use c10::impl::GenericDict(keyType, valueType) instead.");
|
| 73 |
+
static_assert(!std::is_same_v<Value, IValue>, "This constructor is not valid for Dict<_, IValue>. Please use c10::impl::GenericDict(keyType, valueType) instead.");
|
| 74 |
+
}
|
| 75 |
+
|
| 76 |
+
template<class Key, class Value>
|
| 77 |
+
Dict<Key, Value>::Dict(TypePtr keyType, TypePtr valueType)
|
| 78 |
+
: Dict(make_intrusive<detail::DictImpl>(
|
| 79 |
+
detail::DictImpl::dict_map_type(),
|
| 80 |
+
detail::DictImpl::DictElementTypes {std::move(keyType), std::move(valueType)})) {
|
| 81 |
+
static_assert(std::is_same_v<Key, IValue>, "This constructor is only valid for c10::impl::GenericDict.");
|
| 82 |
+
static_assert(std::is_same_v<Value, IValue>, "This constructor is only valid for c10::impl::GenericDict.");
|
| 83 |
+
}
|
| 84 |
+
|
| 85 |
+
template<class Key, class Value>
|
| 86 |
+
Dict<Key, Value>::Dict(c10::intrusive_ptr<detail::DictImpl>&& impl): impl_(std::move(impl)) {}
|
| 87 |
+
|
| 88 |
+
template<class Key, class Value>
|
| 89 |
+
Dict<Key, Value> Dict<Key, Value>::copy() const {
|
| 90 |
+
return Dict<Key, Value>(impl_->copy());
|
| 91 |
+
}
|
| 92 |
+
|
| 93 |
+
template<class Key, class Value>
|
| 94 |
+
typename Dict<Key, Value>::iterator Dict<Key, Value>::begin() const {
|
| 95 |
+
return iterator{impl_->dict.begin()};
|
| 96 |
+
}
|
| 97 |
+
|
| 98 |
+
template<class Key, class Value>
|
| 99 |
+
typename Dict<Key, Value>::iterator Dict<Key, Value>::end() const {
|
| 100 |
+
return iterator{impl_->dict.end()};
|
| 101 |
+
}
|
| 102 |
+
|
| 103 |
+
template<class Key, class Value>
|
| 104 |
+
bool Dict<Key, Value>::empty() const {
|
| 105 |
+
return impl_->dict.empty();
|
| 106 |
+
}
|
| 107 |
+
|
| 108 |
+
template<class Key, class Value>
|
| 109 |
+
typename Dict<Key, Value>::size_type Dict<Key, Value>::size() const {
|
| 110 |
+
return impl_->dict.size();
|
| 111 |
+
}
|
| 112 |
+
|
| 113 |
+
template<class Key, class Value>
|
| 114 |
+
void Dict<Key, Value>::clear() const {
|
| 115 |
+
impl_->dict.clear();
|
| 116 |
+
}
|
| 117 |
+
|
| 118 |
+
template<class Key, class Value>
|
| 119 |
+
template<class Key_, class Value_>
|
| 120 |
+
std::pair<typename Dict<Key, Value>::iterator, bool> Dict<Key, Value>::insert(Key_&& key, Value_&& value) const {
|
| 121 |
+
static_assert(std::is_constructible_v<Key, Key_>, "Wrong type for the key argument of Dict::insert");
|
| 122 |
+
static_assert(std::is_constructible_v<Value, Value_>, "Wrong type for the value argument of Dict::insert");
|
| 123 |
+
auto inserted = impl_->dict.emplace(
|
| 124 |
+
Key(std::forward<Key_>(key)),
|
| 125 |
+
Value(std::forward<Value_>(value)));
|
| 126 |
+
return {iterator{inserted.first}, inserted.second};
|
| 127 |
+
}
|
| 128 |
+
|
| 129 |
+
template<class Key, class Value>
|
| 130 |
+
template<class Key_, class Value_>
|
| 131 |
+
std::pair<typename Dict<Key, Value>::iterator, bool> Dict<Key, Value>::insert_or_assign(Key_&& key, Value_&& value) const {
|
| 132 |
+
static_assert(std::is_constructible_v<Key, Key_>, "Wrong type for the key argument of Dict::insert_or_assign");
|
| 133 |
+
static_assert(std::is_constructible_v<Value, Value_>, "Wrong type for the value argument of Dict::insert_or_assign");
|
| 134 |
+
auto inserted = impl_->dict.insert_or_assign(
|
| 135 |
+
Key(std::forward<Key_>(key)),
|
| 136 |
+
Value(std::forward<Value_>(value)));
|
| 137 |
+
return {iterator{inserted.first}, inserted.second};
|
| 138 |
+
}
|
| 139 |
+
|
| 140 |
+
template<class Key, class Value>
|
| 141 |
+
void Dict<Key, Value>::erase(iterator iter) const {
|
| 142 |
+
impl_->dict.erase(iter.entryRef_.iterator_);
|
| 143 |
+
}
|
| 144 |
+
|
| 145 |
+
template <class Key, class Value>
|
| 146 |
+
[[nodiscard]] size_t Dict<Key, Value>::erase(const Key& key) const {
|
| 147 |
+
return impl_->dict.erase(key);
|
| 148 |
+
}
|
| 149 |
+
|
| 150 |
+
template<class Key, class Value>
|
| 151 |
+
Value Dict<Key, Value>::at(const Key& key) const {
|
| 152 |
+
return impl_->dict.at(key).template to<Value>();
|
| 153 |
+
}
|
| 154 |
+
|
| 155 |
+
template<class Key, class Value>
|
| 156 |
+
typename Dict<Key, Value>::iterator Dict<Key, Value>::find(const Key& key) const {
|
| 157 |
+
return iterator{impl_->dict.find(key)};
|
| 158 |
+
}
|
| 159 |
+
|
| 160 |
+
template<class Key, class Value>
|
| 161 |
+
bool Dict<Key, Value>::contains(const Key& key) const {
|
| 162 |
+
return end() != find(key);
|
| 163 |
+
}
|
| 164 |
+
|
| 165 |
+
template<class Key, class Value>
|
| 166 |
+
void Dict<Key, Value>::reserve(size_type count) const {
|
| 167 |
+
impl_->dict.reserve(count);
|
| 168 |
+
}
|
| 169 |
+
|
| 170 |
+
template<class Key, class Value>
|
| 171 |
+
TypePtr Dict<Key, Value>::keyType() const {
|
| 172 |
+
return impl_->elementTypes.keyType;
|
| 173 |
+
}
|
| 174 |
+
|
| 175 |
+
template<class Key, class Value>
|
| 176 |
+
TypePtr Dict<Key, Value>::valueType() const {
|
| 177 |
+
return impl_->elementTypes.valueType;
|
| 178 |
+
}
|
| 179 |
+
template <class Key, class Value>
|
| 180 |
+
void Dict<Key, Value>::unsafeSetKeyType(TypePtr t) {
|
| 181 |
+
impl_->elementTypes.keyType = std::move(t);
|
| 182 |
+
}
|
| 183 |
+
|
| 184 |
+
template <class Key, class Value>
|
| 185 |
+
void Dict<Key, Value>::unsafeSetValueType(TypePtr t) {
|
| 186 |
+
impl_->elementTypes.valueType = std::move(t);
|
| 187 |
+
}
|
| 188 |
+
|
| 189 |
+
template <class Key_, class Value_>
|
| 190 |
+
bool operator==(const Dict<Key_, Value_>& lhs, const Dict<Key_, Value_>& rhs) {
|
| 191 |
+
// Dicts with the same identity trivially compare equal.
|
| 192 |
+
if (lhs.impl_ == rhs.impl_) {
|
| 193 |
+
return true;
|
| 194 |
+
}
|
| 195 |
+
|
| 196 |
+
// Otherwise compare the values
|
| 197 |
+
return *lhs.impl_ == *rhs.impl_;
|
| 198 |
+
}
|
| 199 |
+
|
| 200 |
+
template <class Key_, class Value_>
|
| 201 |
+
bool operator!=(const Dict<Key_, Value_>& lhs, const Dict<Key_, Value_>& rhs) {
|
| 202 |
+
return !(lhs == rhs);
|
| 203 |
+
}
|
| 204 |
+
|
| 205 |
+
template <class Key, class Value>
|
| 206 |
+
bool Dict<Key, Value>::is(const Dict& rhs) const {
|
| 207 |
+
return this->impl_ == rhs.impl_;
|
| 208 |
+
}
|
| 209 |
+
}
|
lib/python3.12/site-packages/torch/include/ATen/core/DimVector.h
ADDED
|
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
#include <c10/util/DimVector.h>
|
| 3 |
+
|
| 4 |
+
namespace at {
|
| 5 |
+
|
| 6 |
+
// Re-declaring 'DimVector' type and size inside 'at' namespace.
|
| 7 |
+
// This is done to avoid modifying every use into their 'c10'
|
| 8 |
+
// equivalent.
|
| 9 |
+
|
| 10 |
+
using c10::kDimVectorStaticSize;
|
| 11 |
+
using c10::DimVector;
|
| 12 |
+
|
| 13 |
+
} // namespace at
|
lib/python3.12/site-packages/torch/include/ATen/core/Dimname.h
ADDED
|
@@ -0,0 +1,48 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <ATen/core/symbol.h>
|
| 4 |
+
#include <c10/util/ArrayRef.h>
|
| 5 |
+
#include <optional>
|
| 6 |
+
#include <ostream>
|
| 7 |
+
|
| 8 |
+
namespace at {
|
| 9 |
+
|
| 10 |
+
enum class NameType: uint8_t { BASIC, WILDCARD };
|
| 11 |
+
|
| 12 |
+
struct TORCH_API Dimname {
|
| 13 |
+
static Dimname fromSymbol(Symbol name);
|
| 14 |
+
static Dimname wildcard();
|
| 15 |
+
static bool isValidName(const std::string& name);
|
| 16 |
+
|
| 17 |
+
NameType type() const { return type_; }
|
| 18 |
+
Symbol symbol() const { return name_; }
|
| 19 |
+
|
| 20 |
+
bool isBasic() const { return type_ == NameType::BASIC; }
|
| 21 |
+
bool isWildcard() const { return type_ == NameType::WILDCARD; }
|
| 22 |
+
|
| 23 |
+
bool matches(Dimname other) const;
|
| 24 |
+
std::optional<Dimname> unify(Dimname other) const;
|
| 25 |
+
|
| 26 |
+
private:
|
| 27 |
+
Dimname(Symbol name)
|
| 28 |
+
: name_(name), type_(NameType::BASIC) {}
|
| 29 |
+
Dimname(Symbol name, NameType type)
|
| 30 |
+
: name_(name), type_(type) {}
|
| 31 |
+
|
| 32 |
+
Symbol name_;
|
| 33 |
+
NameType type_;
|
| 34 |
+
};
|
| 35 |
+
|
| 36 |
+
using DimnameList = c10::ArrayRef<Dimname>;
|
| 37 |
+
|
| 38 |
+
TORCH_API std::ostream& operator<<(std::ostream& out, const Dimname& dimname);
|
| 39 |
+
|
| 40 |
+
inline bool operator==(const Dimname& lhs, const Dimname& rhs) {
|
| 41 |
+
return lhs.symbol() == rhs.symbol();
|
| 42 |
+
}
|
| 43 |
+
|
| 44 |
+
inline bool operator!=(const Dimname& lhs, const Dimname& rhs) {
|
| 45 |
+
return !(lhs == rhs);
|
| 46 |
+
}
|
| 47 |
+
|
| 48 |
+
} // namespace at
|
lib/python3.12/site-packages/torch/include/ATen/core/DistributionsHelper.h
ADDED
|
@@ -0,0 +1,329 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <ATen/core/Array.h>
|
| 4 |
+
#include <ATen/core/TransformationHelper.h>
|
| 5 |
+
#include <c10/util/Half.h>
|
| 6 |
+
#include <c10/util/BFloat16.h>
|
| 7 |
+
#include <c10/util/MathConstants.h>
|
| 8 |
+
#include <c10/macros/Macros.h>
|
| 9 |
+
|
| 10 |
+
#include <cmath>
|
| 11 |
+
#include <limits>
|
| 12 |
+
#include <optional>
|
| 13 |
+
#include <type_traits>
|
| 14 |
+
|
| 15 |
+
/**
|
| 16 |
+
* Distributions kernel adapted from THRandom.cpp
|
| 17 |
+
* The kernels try to follow std::random distributions signature
|
| 18 |
+
* For instance: in ATen
|
| 19 |
+
* auto gen = at::detail::createCPUGenerator();
|
| 20 |
+
* at::uniform_real_distribution<double> uniform(0, 1);
|
| 21 |
+
* auto sample = uniform(gen.get());
|
| 22 |
+
*
|
| 23 |
+
* vs std::random
|
| 24 |
+
*
|
| 25 |
+
* std::mt19937 gen;
|
| 26 |
+
* std::uniform_real_distribution uniform(0, 1);
|
| 27 |
+
* auto sample = uniform(gen);
|
| 28 |
+
*/
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
namespace at {
|
| 32 |
+
namespace {
|
| 33 |
+
|
| 34 |
+
/**
|
| 35 |
+
* Samples a discrete uniform distribution in the range [base, base+range) of type T
|
| 36 |
+
*/
|
| 37 |
+
template <typename T>
|
| 38 |
+
struct uniform_int_from_to_distribution {
|
| 39 |
+
|
| 40 |
+
C10_HOST_DEVICE inline uniform_int_from_to_distribution(uint64_t range, int64_t base) : range_(range), base_(base) {}
|
| 41 |
+
|
| 42 |
+
template <typename RNG>
|
| 43 |
+
C10_HOST_DEVICE inline T operator()(RNG generator) {
|
| 44 |
+
if ((
|
| 45 |
+
std::is_same_v<T, int64_t> ||
|
| 46 |
+
std::is_same_v<T, double> ||
|
| 47 |
+
std::is_same_v<T, float> ||
|
| 48 |
+
std::is_same_v<T, at::BFloat16>) && range_ >= 1ULL << 32)
|
| 49 |
+
{
|
| 50 |
+
return transformation::uniform_int_from_to<T>(generator->random64(), range_, base_);
|
| 51 |
+
} else {
|
| 52 |
+
return transformation::uniform_int_from_to<T>(generator->random(), range_, base_);
|
| 53 |
+
}
|
| 54 |
+
}
|
| 55 |
+
|
| 56 |
+
private:
|
| 57 |
+
uint64_t range_;
|
| 58 |
+
int64_t base_;
|
| 59 |
+
};
|
| 60 |
+
|
| 61 |
+
/**
|
| 62 |
+
* Samples a discrete uniform distribution in the range [min_value(int64_t), max_value(int64_t)]
|
| 63 |
+
*/
|
| 64 |
+
template <typename T>
|
| 65 |
+
struct uniform_int_full_range_distribution {
|
| 66 |
+
|
| 67 |
+
template <typename RNG>
|
| 68 |
+
C10_HOST_DEVICE inline T operator()(RNG generator) {
|
| 69 |
+
return transformation::uniform_int_full_range<T>(generator->random64());
|
| 70 |
+
}
|
| 71 |
+
|
| 72 |
+
};
|
| 73 |
+
|
| 74 |
+
/**
|
| 75 |
+
* Samples a discrete uniform distribution in the range [0, max_value(T)] for integral types
|
| 76 |
+
* and [0, 2^mantissa] for floating-point types.
|
| 77 |
+
*/
|
| 78 |
+
template <typename T>
|
| 79 |
+
struct uniform_int_distribution {
|
| 80 |
+
|
| 81 |
+
template <typename RNG>
|
| 82 |
+
C10_HOST_DEVICE inline T operator()(RNG generator) {
|
| 83 |
+
if constexpr (std::is_same_v<T, double> || std::is_same_v<T, int64_t>) {
|
| 84 |
+
return transformation::uniform_int<T>(generator->random64());
|
| 85 |
+
} else {
|
| 86 |
+
return transformation::uniform_int<T>(generator->random());
|
| 87 |
+
}
|
| 88 |
+
}
|
| 89 |
+
|
| 90 |
+
};
|
| 91 |
+
|
| 92 |
+
/**
|
| 93 |
+
* Samples a uniform distribution in the range [from, to) of type T
|
| 94 |
+
*/
|
| 95 |
+
template <typename T>
|
| 96 |
+
struct uniform_real_distribution {
|
| 97 |
+
|
| 98 |
+
C10_HOST_DEVICE inline uniform_real_distribution(T from, T to) : from_(from), to_(to) {
|
| 99 |
+
TORCH_CHECK_IF_NOT_ON_CUDA(from <= to);
|
| 100 |
+
TORCH_CHECK_IF_NOT_ON_CUDA(to - from <= std::numeric_limits<T>::max());
|
| 101 |
+
}
|
| 102 |
+
|
| 103 |
+
template <typename RNG>
|
| 104 |
+
C10_HOST_DEVICE inline dist_acctype<T> operator()(RNG generator){
|
| 105 |
+
if constexpr (std::is_same_v<T, double>) {
|
| 106 |
+
return transformation::uniform_real<T>(generator->random64(), from_, to_);
|
| 107 |
+
} else {
|
| 108 |
+
return transformation::uniform_real<T>(generator->random(), from_, to_);
|
| 109 |
+
}
|
| 110 |
+
}
|
| 111 |
+
|
| 112 |
+
private:
|
| 113 |
+
T from_;
|
| 114 |
+
T to_;
|
| 115 |
+
};
|
| 116 |
+
|
| 117 |
+
// The SFINAE checks introduced in #39816 looks overcomplicated and must revisited
|
| 118 |
+
// https://github.com/pytorch/pytorch/issues/40052
|
| 119 |
+
#define DISTRIBUTION_HELPER_GENERATE_HAS_MEMBER(member) \
|
| 120 |
+
template <typename T> \
|
| 121 |
+
struct has_member_##member \
|
| 122 |
+
{ \
|
| 123 |
+
typedef char yes; \
|
| 124 |
+
typedef long no; \
|
| 125 |
+
template <typename U> static yes test(decltype(&U::member)); \
|
| 126 |
+
template <typename U> static no test(...); \
|
| 127 |
+
static constexpr bool value = sizeof(test<T>(0)) == sizeof(yes); \
|
| 128 |
+
}
|
| 129 |
+
|
| 130 |
+
DISTRIBUTION_HELPER_GENERATE_HAS_MEMBER(next_double_normal_sample);
|
| 131 |
+
DISTRIBUTION_HELPER_GENERATE_HAS_MEMBER(set_next_double_normal_sample);
|
| 132 |
+
DISTRIBUTION_HELPER_GENERATE_HAS_MEMBER(next_float_normal_sample);
|
| 133 |
+
DISTRIBUTION_HELPER_GENERATE_HAS_MEMBER(set_next_float_normal_sample);
|
| 134 |
+
|
| 135 |
+
#define DISTRIBUTION_HELPER_GENERATE_NEXT_NORMAL_METHODS(TYPE) \
|
| 136 |
+
\
|
| 137 |
+
template <typename RNG, typename ret_type, \
|
| 138 |
+
typename std::enable_if_t<( \
|
| 139 |
+
has_member_next_##TYPE##_normal_sample<RNG>::value && \
|
| 140 |
+
has_member_set_next_##TYPE##_normal_sample<RNG>::value \
|
| 141 |
+
), int> = 0> \
|
| 142 |
+
C10_HOST_DEVICE inline bool maybe_get_next_##TYPE##_normal_sample(RNG* generator, ret_type* ret) { \
|
| 143 |
+
if (generator->next_##TYPE##_normal_sample()) { \
|
| 144 |
+
*ret = *(generator->next_##TYPE##_normal_sample()); \
|
| 145 |
+
generator->set_next_##TYPE##_normal_sample(std::optional<TYPE>()); \
|
| 146 |
+
return true; \
|
| 147 |
+
} \
|
| 148 |
+
return false; \
|
| 149 |
+
} \
|
| 150 |
+
\
|
| 151 |
+
template <typename RNG, typename ret_type, \
|
| 152 |
+
typename std::enable_if_t<( \
|
| 153 |
+
!has_member_next_##TYPE##_normal_sample<RNG>::value || \
|
| 154 |
+
!has_member_set_next_##TYPE##_normal_sample<RNG>::value \
|
| 155 |
+
), int> = 0> \
|
| 156 |
+
C10_HOST_DEVICE inline bool maybe_get_next_##TYPE##_normal_sample(RNG* /*generator*/, ret_type* /*ret*/) { \
|
| 157 |
+
return false; \
|
| 158 |
+
} \
|
| 159 |
+
\
|
| 160 |
+
template <typename RNG, typename ret_type, \
|
| 161 |
+
typename std::enable_if_t<( \
|
| 162 |
+
has_member_set_next_##TYPE##_normal_sample<RNG>::value \
|
| 163 |
+
), int> = 0> \
|
| 164 |
+
C10_HOST_DEVICE inline void maybe_set_next_##TYPE##_normal_sample(RNG* generator, ret_type cache) { \
|
| 165 |
+
generator->set_next_##TYPE##_normal_sample(cache); \
|
| 166 |
+
} \
|
| 167 |
+
\
|
| 168 |
+
template <typename RNG, typename ret_type, \
|
| 169 |
+
typename std::enable_if_t<( \
|
| 170 |
+
!has_member_set_next_##TYPE##_normal_sample<RNG>::value \
|
| 171 |
+
), int> = 0> \
|
| 172 |
+
C10_HOST_DEVICE inline void maybe_set_next_##TYPE##_normal_sample(RNG* /*generator*/, ret_type /*cache*/) { \
|
| 173 |
+
}
|
| 174 |
+
|
| 175 |
+
DISTRIBUTION_HELPER_GENERATE_NEXT_NORMAL_METHODS(double)
|
| 176 |
+
DISTRIBUTION_HELPER_GENERATE_NEXT_NORMAL_METHODS(float)
|
| 177 |
+
|
| 178 |
+
/**
|
| 179 |
+
* Samples a normal distribution using the Box-Muller method
|
| 180 |
+
* Takes mean and standard deviation as inputs
|
| 181 |
+
* Note that Box-muller method returns two samples at a time.
|
| 182 |
+
* Hence, we cache the "next" sample in the CPUGeneratorImpl class.
|
| 183 |
+
*/
|
| 184 |
+
template <typename T>
|
| 185 |
+
struct normal_distribution {
|
| 186 |
+
|
| 187 |
+
C10_HOST_DEVICE inline normal_distribution(T mean_in, T stdv_in) : mean(mean_in), stdv(stdv_in) {
|
| 188 |
+
TORCH_CHECK_IF_NOT_ON_CUDA(stdv_in >= 0, "stdv_in must be positive: ", stdv_in);
|
| 189 |
+
}
|
| 190 |
+
|
| 191 |
+
template <typename RNG>
|
| 192 |
+
C10_HOST_DEVICE inline dist_acctype<T> operator()(RNG generator){
|
| 193 |
+
dist_acctype<T> ret;
|
| 194 |
+
// return cached values if available
|
| 195 |
+
if constexpr (std::is_same_v<T, double>) {
|
| 196 |
+
if (maybe_get_next_double_normal_sample(generator, &ret)) {
|
| 197 |
+
return transformation::normal(ret, mean, stdv);
|
| 198 |
+
}
|
| 199 |
+
} else {
|
| 200 |
+
if (maybe_get_next_float_normal_sample(generator, &ret)) {
|
| 201 |
+
return transformation::normal(ret, mean, stdv);
|
| 202 |
+
}
|
| 203 |
+
}
|
| 204 |
+
// otherwise generate new normal values
|
| 205 |
+
uniform_real_distribution<T> uniform(0.0, 1.0);
|
| 206 |
+
const dist_acctype<T> u1 = uniform(generator);
|
| 207 |
+
const dist_acctype<T> u2 = uniform(generator);
|
| 208 |
+
const dist_acctype<T> r = ::sqrt(static_cast<T>(-2.0) * ::log1p(-u2));
|
| 209 |
+
const dist_acctype<T> theta = static_cast<T>(2.0) * c10::pi<T> * u1;
|
| 210 |
+
if constexpr (std::is_same_v<T, double>) {
|
| 211 |
+
maybe_set_next_double_normal_sample(generator, r * ::sin(theta));
|
| 212 |
+
} else {
|
| 213 |
+
maybe_set_next_float_normal_sample(generator, r * ::sin(theta));
|
| 214 |
+
}
|
| 215 |
+
ret = r * ::cos(theta);
|
| 216 |
+
return transformation::normal(ret, mean, stdv);
|
| 217 |
+
}
|
| 218 |
+
|
| 219 |
+
private:
|
| 220 |
+
T mean;
|
| 221 |
+
T stdv;
|
| 222 |
+
};
|
| 223 |
+
|
| 224 |
+
template <typename T>
|
| 225 |
+
struct DiscreteDistributionType { using type = float; };
|
| 226 |
+
|
| 227 |
+
template <> struct DiscreteDistributionType<double> { using type = double; };
|
| 228 |
+
|
| 229 |
+
/**
|
| 230 |
+
* Samples a bernoulli distribution given a probability input
|
| 231 |
+
*/
|
| 232 |
+
template <typename T>
|
| 233 |
+
struct bernoulli_distribution {
|
| 234 |
+
|
| 235 |
+
C10_HOST_DEVICE inline bernoulli_distribution(T p_in) : p(p_in) {
|
| 236 |
+
TORCH_CHECK_IF_NOT_ON_CUDA(p_in >= 0 && p_in <= 1);
|
| 237 |
+
}
|
| 238 |
+
|
| 239 |
+
template <typename RNG>
|
| 240 |
+
C10_HOST_DEVICE inline T operator()(RNG generator) {
|
| 241 |
+
uniform_real_distribution<T> uniform(0.0, 1.0);
|
| 242 |
+
return transformation::bernoulli<T>(uniform(generator), p);
|
| 243 |
+
}
|
| 244 |
+
|
| 245 |
+
private:
|
| 246 |
+
T p;
|
| 247 |
+
};
|
| 248 |
+
|
| 249 |
+
/**
|
| 250 |
+
* Samples a geometric distribution given a probability input
|
| 251 |
+
*/
|
| 252 |
+
template <typename T>
|
| 253 |
+
struct geometric_distribution {
|
| 254 |
+
|
| 255 |
+
C10_HOST_DEVICE inline geometric_distribution(T p_in) : p(p_in) {
|
| 256 |
+
TORCH_CHECK_IF_NOT_ON_CUDA(p_in > 0 && p_in < 1);
|
| 257 |
+
}
|
| 258 |
+
|
| 259 |
+
template <typename RNG>
|
| 260 |
+
C10_HOST_DEVICE inline T operator()(RNG generator) {
|
| 261 |
+
uniform_real_distribution<T> uniform(0.0, 1.0);
|
| 262 |
+
return transformation::geometric<T>(uniform(generator), p);
|
| 263 |
+
}
|
| 264 |
+
|
| 265 |
+
private:
|
| 266 |
+
T p;
|
| 267 |
+
};
|
| 268 |
+
|
| 269 |
+
/**
|
| 270 |
+
* Samples an exponential distribution given a lambda input
|
| 271 |
+
*/
|
| 272 |
+
template <typename T>
|
| 273 |
+
struct exponential_distribution {
|
| 274 |
+
|
| 275 |
+
C10_HOST_DEVICE inline exponential_distribution(T lambda_in) : lambda(lambda_in) {}
|
| 276 |
+
|
| 277 |
+
template <typename RNG>
|
| 278 |
+
C10_HOST_DEVICE inline T operator()(RNG generator) {
|
| 279 |
+
uniform_real_distribution<T> uniform(0.0, 1.0);
|
| 280 |
+
return transformation::exponential<T>(uniform(generator), lambda);
|
| 281 |
+
}
|
| 282 |
+
|
| 283 |
+
private:
|
| 284 |
+
T lambda;
|
| 285 |
+
};
|
| 286 |
+
|
| 287 |
+
/**
|
| 288 |
+
* Samples a cauchy distribution given median and sigma as inputs
|
| 289 |
+
*/
|
| 290 |
+
template <typename T>
|
| 291 |
+
struct cauchy_distribution {
|
| 292 |
+
|
| 293 |
+
C10_HOST_DEVICE inline cauchy_distribution(T median_in, T sigma_in) : median(median_in), sigma(sigma_in) {}
|
| 294 |
+
|
| 295 |
+
template <typename RNG>
|
| 296 |
+
C10_HOST_DEVICE inline T operator()(RNG generator) {
|
| 297 |
+
uniform_real_distribution<T> uniform(0.0, 1.0);
|
| 298 |
+
return transformation::cauchy<T>(uniform(generator), median, sigma);
|
| 299 |
+
}
|
| 300 |
+
|
| 301 |
+
private:
|
| 302 |
+
T median;
|
| 303 |
+
T sigma;
|
| 304 |
+
};
|
| 305 |
+
|
| 306 |
+
/**
|
| 307 |
+
* Samples a lognormal distribution
|
| 308 |
+
* Takes mean and standard deviation as inputs
|
| 309 |
+
* Outputs two samples at a time
|
| 310 |
+
*/
|
| 311 |
+
template <typename T>
|
| 312 |
+
struct lognormal_distribution {
|
| 313 |
+
|
| 314 |
+
C10_HOST_DEVICE inline lognormal_distribution(T mean_in, T stdv_in) : mean(mean_in), stdv(stdv_in) {
|
| 315 |
+
TORCH_CHECK_IF_NOT_ON_CUDA(stdv_in > 0);
|
| 316 |
+
}
|
| 317 |
+
|
| 318 |
+
template<typename RNG>
|
| 319 |
+
C10_HOST_DEVICE inline T operator()(RNG generator){
|
| 320 |
+
normal_distribution<T> normal(mean, stdv);
|
| 321 |
+
return transformation::log_normal<T>(normal(generator));
|
| 322 |
+
}
|
| 323 |
+
|
| 324 |
+
private:
|
| 325 |
+
T mean;
|
| 326 |
+
T stdv;
|
| 327 |
+
};
|
| 328 |
+
}
|
| 329 |
+
} // namespace at
|
lib/python3.12/site-packages/torch/include/ATen/core/Formatting.h
ADDED
|
@@ -0,0 +1,25 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <ostream>
|
| 4 |
+
#include <string>
|
| 5 |
+
|
| 6 |
+
#include <c10/core/Scalar.h>
|
| 7 |
+
#include <ATen/core/Tensor.h>
|
| 8 |
+
|
| 9 |
+
namespace c10 {
|
| 10 |
+
TORCH_API std::ostream& operator<<(std::ostream& out, Backend b);
|
| 11 |
+
TORCH_API std::ostream& operator<<(std::ostream & out, const Scalar& s);
|
| 12 |
+
TORCH_API std::string toString(const Scalar& s);
|
| 13 |
+
}
|
| 14 |
+
namespace at {
|
| 15 |
+
|
| 16 |
+
TORCH_API std::ostream& operator<<(std::ostream& out, const DeprecatedTypeProperties& t);
|
| 17 |
+
TORCH_API std::ostream& print(
|
| 18 |
+
std::ostream& stream,
|
| 19 |
+
const Tensor& tensor,
|
| 20 |
+
int64_t linesize);
|
| 21 |
+
inline std::ostream& operator<<(std::ostream & out, const Tensor & t) {
|
| 22 |
+
return print(out,t,80);
|
| 23 |
+
}
|
| 24 |
+
TORCH_API void print(const Tensor & t, int64_t linesize=80);
|
| 25 |
+
}
|
lib/python3.12/site-packages/torch/include/ATen/core/Generator.h
ADDED
|
@@ -0,0 +1,191 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#pragma once
|
| 2 |
+
|
| 3 |
+
#include <cstdint>
|
| 4 |
+
#include <deque>
|
| 5 |
+
#include <mutex>
|
| 6 |
+
#include <utility>
|
| 7 |
+
|
| 8 |
+
#include <c10/util/Exception.h>
|
| 9 |
+
#include <c10/util/intrusive_ptr.h>
|
| 10 |
+
#include <c10/core/Device.h>
|
| 11 |
+
#include <c10/core/DispatchKeySet.h>
|
| 12 |
+
|
| 13 |
+
// For the record I don't think this is a correct pimpl idiom.
|
| 14 |
+
// Including Impl header in interface header defeats the purpose
|
| 15 |
+
// because you can't change Impl private members without forcing
|
| 16 |
+
// everything that included the interface to rebuild.
|
| 17 |
+
// Impl should be forward-declared in the interface header instead.
|
| 18 |
+
#include <c10/core/GeneratorImpl.h>
|
| 19 |
+
|
| 20 |
+
/**
|
| 21 |
+
* Note [Generator]
|
| 22 |
+
* ~~~~~~~~~~~~~~~~
|
| 23 |
+
* A Pseudo Random Number Generator (PRNG) is an engine that uses an algorithm to
|
| 24 |
+
* generate a seemingly random sequence of numbers, that may be later be used in creating
|
| 25 |
+
* a random distribution. Such an engine almost always maintains a state and requires a
|
| 26 |
+
* seed to start off the creation of random numbers. Often times, users have
|
| 27 |
+
* found it beneficial to be able to explicitly create, retain, and destroy
|
| 28 |
+
* PRNG states and also be able to have control over the seed value.
|
| 29 |
+
*
|
| 30 |
+
* A Generator in ATen gives users the ability to read, write and modify a PRNG engine.
|
| 31 |
+
* For instance, it does so by letting users seed a PRNG engine, fork the state of the
|
| 32 |
+
* engine, etc.
|
| 33 |
+
*
|
| 34 |
+
* By default, there is one generator per device, and a device's generator is
|
| 35 |
+
* lazily created. A user can use the torch.Generator() api to create their own generator.
|
| 36 |
+
*/
|
| 37 |
+
|
| 38 |
+
/**
|
| 39 |
+
* Note [Acquire lock when using random generators]
|
| 40 |
+
* ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
| 41 |
+
* Generator and its derived classes are NOT thread-safe. Please note that most of the
|
| 42 |
+
* places where we have inserted locking for generators are historically based, and we
|
| 43 |
+
* haven't actually checked that everything is truly thread safe (and it probably isn't).
|
| 44 |
+
* Please use the public mutex_ when using any methods from these classes, except for the
|
| 45 |
+
* read-only methods. You can learn about the usage by looking into the unittests
|
| 46 |
+
* (aten/src/ATen/cpu_generator_test.cpp) and other places where we have used lock_guard.
|
| 47 |
+
*
|
| 48 |
+
* TODO: Look into changing the threading semantics of Generators in ATen (e.g., making
|
| 49 |
+
* them non-thread safe and instead making the generator state splittable, to accommodate
|
| 50 |
+
* forks into other threads).
|
| 51 |
+
*/
|
| 52 |
+
|
| 53 |
+
namespace at {
|
| 54 |
+
|
| 55 |
+
class Tensor;
|
| 56 |
+
|
| 57 |
+
struct TORCH_API Generator {
|
| 58 |
+
Generator() = default;
|
| 59 |
+
|
| 60 |
+
explicit Generator(c10::intrusive_ptr<c10::GeneratorImpl> gen_impl)
|
| 61 |
+
: impl_(std::move(gen_impl)) {
|
| 62 |
+
if (impl_.get() == nullptr) {
|
| 63 |
+
throw std::runtime_error("GeneratorImpl with nullptr is not supported");
|
| 64 |
+
}
|
| 65 |
+
}
|
| 66 |
+
|
| 67 |
+
bool operator==(const Generator& rhs) const {
|
| 68 |
+
return this->impl_ == rhs.impl_;
|
| 69 |
+
}
|
| 70 |
+
|
| 71 |
+
bool operator!=(const Generator& rhs) const {
|
| 72 |
+
return !((*this) == rhs);
|
| 73 |
+
}
|
| 74 |
+
|
| 75 |
+
bool defined() const {
|
| 76 |
+
return static_cast<bool>(impl_);
|
| 77 |
+
}
|
| 78 |
+
|
| 79 |
+
c10::GeneratorImpl* unsafeGetGeneratorImpl() const {
|
| 80 |
+
return impl_.get();
|
| 81 |
+
}
|
| 82 |
+
|
| 83 |
+
c10::GeneratorImpl* unsafeReleaseGeneratorImpl() {
|
| 84 |
+
return impl_.release();
|
| 85 |
+
}
|
| 86 |
+
|
| 87 |
+
const c10::intrusive_ptr<c10::GeneratorImpl>& getIntrusivePtr() const {
|
| 88 |
+
return impl_;
|
| 89 |
+
}
|
| 90 |
+
|
| 91 |
+
void set_current_seed(uint64_t seed) { impl_->set_current_seed(seed); }
|
| 92 |
+
// Sets the offset of Generator state to the desired offset. This is currently
|
| 93 |
+
// supported for only Philox based Generators, i.e., CUDA and MPS.
|
| 94 |
+
void set_offset(uint64_t offset) { impl_->set_offset(offset); }
|
| 95 |
+
|
| 96 |
+
// Returns the offset of Generator state. This is currently supported for only
|
| 97 |
+
// Philox based Generators, i.e., CUDA and MPS.
|
| 98 |
+
uint64_t get_offset() const { return impl_->get_offset(); }
|
| 99 |
+
|
| 100 |
+
uint64_t current_seed() const { return impl_->current_seed(); }
|
| 101 |
+
|
| 102 |
+
uint64_t seed() { return impl_->seed(); }
|
| 103 |
+
|
| 104 |
+
// Implementation not inlined to prevent cycle reference between
|
| 105 |
+
// `ATen/core/Generator.h` and `ATen/core/Tensor.h`
|
| 106 |
+
void set_state(const at::Tensor& new_state);
|
| 107 |
+
|
| 108 |
+
at::Tensor get_state() const;
|
| 109 |
+
|
| 110 |
+
void graphsafe_set_state(const Generator& new_state);
|
| 111 |
+
|
| 112 |
+
Generator graphsafe_get_state() const;
|
| 113 |
+
|
| 114 |
+
std::mutex& mutex() {
|
| 115 |
+
return impl_->mutex_;
|
| 116 |
+
}
|
| 117 |
+
|
| 118 |
+
DispatchKeySet key_set() const {
|
| 119 |
+
return impl_->key_set();
|
| 120 |
+
}
|
| 121 |
+
|
| 122 |
+
Device device() const { return impl_->device(); }
|
| 123 |
+
|
| 124 |
+
inline void set_pyobj(PyObject* pyobj) const noexcept {
|
| 125 |
+
impl_->set_pyobj(pyobj);
|
| 126 |
+
}
|
| 127 |
+
|
| 128 |
+
inline PyObject* pyobj() const noexcept {
|
| 129 |
+
return impl_->pyobj();
|
| 130 |
+
}
|
| 131 |
+
|
| 132 |
+
template<typename T>
|
| 133 |
+
T* get() const { return static_cast<T*>(impl_.get()); }
|
| 134 |
+
|
| 135 |
+
Generator clone() const {
|
| 136 |
+
return Generator(impl_->clone());
|
| 137 |
+
}
|
| 138 |
+
|
| 139 |
+
private:
|
| 140 |
+
c10::intrusive_ptr<c10::GeneratorImpl> impl_;
|
| 141 |
+
};
|
| 142 |
+
|
| 143 |
+
template<class Impl, class... Args>
|
| 144 |
+
Generator make_generator(Args&&... args) {
|
| 145 |
+
return Generator(c10::make_intrusive<Impl>(std::forward<Args>(args)...));
|
| 146 |
+
}
|
| 147 |
+
|
| 148 |
+
/**
|
| 149 |
+
* Utility function to static cast input Generator* to
|
| 150 |
+
* the backend generator type (CPU/CUDAGeneratorImpl etc.)
|
| 151 |
+
*/
|
| 152 |
+
template <typename T>
|
| 153 |
+
inline T * check_generator(std::optional<Generator> gen) {
|
| 154 |
+
TORCH_CHECK(gen.has_value(), "Expected Generator but received nullopt");
|
| 155 |
+
TORCH_CHECK(gen->defined(), "Generator with undefined implementation is not allowed");
|
| 156 |
+
TORCH_CHECK(T::device_type() == gen->device().type(), "Expected a '", T::device_type(), "' device type for generator but found '", gen->device().type(), "'");
|
| 157 |
+
return gen->get<T>();
|
| 158 |
+
}
|
| 159 |
+
|
| 160 |
+
/**
|
| 161 |
+
* Utility function used in tensor implementations, which
|
| 162 |
+
* supplies the default generator to tensors, if an input generator
|
| 163 |
+
* is not supplied. The input Generator* is also static casted to
|
| 164 |
+
* the backend generator type (CPU/CUDAGeneratorImpl etc.)
|
| 165 |
+
*/
|
| 166 |
+
template <typename T>
|
| 167 |
+
inline T* get_generator_or_default(const std::optional<Generator>& gen, const Generator& default_gen) {
|
| 168 |
+
return gen.has_value() && gen->defined() ? check_generator<T>(gen) : check_generator<T>(default_gen);
|
| 169 |
+
}
|
| 170 |
+
|
| 171 |
+
namespace detail {
|
| 172 |
+
|
| 173 |
+
/**
|
| 174 |
+
* Helper function for checking the validity of new random generator
|
| 175 |
+
* state. Right now following conditions are checked:
|
| 176 |
+
*
|
| 177 |
+
* - The new state tensor must be a torch.ByteTensor
|
| 178 |
+
* - Data of the new state tensor must be contiguous
|
| 179 |
+
*/
|
| 180 |
+
inline void check_rng_state(const c10::TensorImpl& new_state) {
|
| 181 |
+
TORCH_CHECK_TYPE(
|
| 182 |
+
new_state.layout() == kStrided && new_state.device().type() == kCPU && new_state.dtype() == kByte,
|
| 183 |
+
"RNG state must be a torch.ByteTensor"
|
| 184 |
+
);
|
| 185 |
+
|
| 186 |
+
TORCH_CHECK(new_state.is_contiguous(), "RNG state must be contiguous");
|
| 187 |
+
}
|
| 188 |
+
|
| 189 |
+
} // namespace detail
|
| 190 |
+
|
| 191 |
+
} // namespace at
|