1 #define TORCH_ASSERT_ONLY_METHOD_OPERATORS
2 #include <ATen/native/TensorFactories.h>
3
4 #include <ATen/core/Tensor.h>
5 #include <ATen/CPUGeneratorImpl.h>
6 #include <ATen/Dispatch.h>
7 #include <ATen/EmptyTensor.h>
8 #include <ATen/ExpandUtils.h>
9 #include <ATen/Parallel.h>
10 #include <ATen/MapAllocator.h>
11 #include <ATen/SparseCsrTensorUtils.h>
12 #include <ATen/TracerMode.h>
13 #include <ATen/TensorOperators.h>
14 #include <ATen/NamedTensorUtils.h>
15 #include <ATen/native/UnaryOps.h>
16 #include <c10/core/ScalarType.h>
17 #include <c10/core/TensorOptions.h>
18 #include <c10/util/Exception.h>
19 #include <c10/util/irange.h>
20 #include <c10/util/MathConstants.h>
21
22 #ifndef AT_PER_OPERATOR_HEADERS
23 #include <ATen/Functions.h>
24 #include <ATen/NativeFunctions.h>
25 #else
26 #include <ATen/ops/_cast_Byte_native.h>
27 #include <ATen/ops/_cast_Char_native.h>
28 #include <ATen/ops/_cast_Double_native.h>
29 #include <ATen/ops/_cast_Float_native.h>
30 #include <ATen/ops/_cast_Half_native.h>
31 #include <ATen/ops/_cast_Int_native.h>
32 #include <ATen/ops/_cast_Long_native.h>
33 #include <ATen/ops/_cast_Short_native.h>
34 #include <ATen/ops/_dim_arange_native.h>
35 #include <ATen/ops/_efficientzerotensor_native.h>
36 #include <ATen/ops/_empty_affine_quantized.h>
37 #include <ATen/ops/_empty_per_channel_affine_quantized.h>
38 #include <ATen/ops/_sparse_compressed_tensor_with_dims_native.h>
39 #include <ATen/ops/arange.h>
40 #include <ATen/ops/arange_native.h>
41 #include <ATen/ops/bartlett_window_native.h>
42 #include <ATen/ops/blackman_window_native.h>
43 #include <ATen/ops/clone_native.h>
44 #include <ATen/ops/complex.h>
45 #include <ATen/ops/complex_native.h>
46 #include <ATen/ops/cumprod.h>
47 #include <ATen/ops/empty.h>
48 #include <ATen/ops/empty_like.h>
49 #include <ATen/ops/empty_like_native.h>
50 #include <ATen/ops/empty_native.h>
51 #include <ATen/ops/empty_permuted_native.h>
52 #include <ATen/ops/empty_strided.h>
53 #include <ATen/ops/empty_strided_native.h>
54 #include <ATen/ops/eye.h>
55 #include <ATen/ops/eye_native.h>
56 #include <ATen/ops/fill_native.h>
57 #include <ATen/ops/flip.h>
58 #include <ATen/ops/from_file_native.h>
59 #include <ATen/ops/full_like_native.h>
60 #include <ATen/ops/full_native.h>
61 #include <ATen/ops/hamming_window_native.h>
62 #include <ATen/ops/hann_window_native.h>
63 #include <ATen/ops/kaiser_window_native.h>
64 #include <ATen/ops/linspace.h>
65 #include <ATen/ops/linspace_native.h>
66 #include <ATen/ops/logspace.h>
67 #include <ATen/ops/logspace_native.h>
68 #include <ATen/ops/new_empty_native.h>
69 #include <ATen/ops/new_empty_strided_native.h>
70 #include <ATen/ops/new_full_native.h>
71 #include <ATen/ops/new_ones_native.h>
72 #include <ATen/ops/new_zeros_native.h>
73 #include <ATen/ops/normal_native.h>
74 #include <ATen/ops/ones.h>
75 #include <ATen/ops/ones_like_native.h>
76 #include <ATen/ops/ones_native.h>
77 #include <ATen/ops/polar.h>
78 #include <ATen/ops/polar_native.h>
79 #include <ATen/ops/promote_types.h>
80 #include <ATen/ops/rand_like_native.h>
81 #include <ATen/ops/rand_native.h>
82 #include <ATen/ops/randint_like_native.h>
83 #include <ATen/ops/randint_native.h>
84 #include <ATen/ops/randn_like_native.h>
85 #include <ATen/ops/randn_native.h>
86 #include <ATen/ops/randperm.h>
87 #include <ATen/ops/randperm_native.h>
88 #include <ATen/ops/range.h>
89 #include <ATen/ops/range_native.h>
90 #include <ATen/ops/scalar_tensor_native.h>
91 #include <ATen/ops/tril_indices_native.h>
92 #include <ATen/ops/triu_indices_native.h>
93 #include <ATen/ops/vander_native.h>
94 #include <ATen/ops/zeros_like_native.h>
95 #include <ATen/ops/zeros_like_ops.h>
96 #include <ATen/ops/zeros_native.h>
97 #endif
98
99 #include <c10/core/SymIntArrayRef.h>
100 #include <algorithm>
101 #include <cstddef>
102 #include <string>
103 #include <utility>
104
105 namespace at::native {
106 namespace {
window_function_checks(const char * function_name,const TensorOptions & options,int64_t window_length)107 void window_function_checks(
108 const char* function_name,
109 const TensorOptions& options,
110 int64_t window_length) {
111 TORCH_CHECK(
112 options.layout() != kSparse,
113 function_name,
114 " is not implemented for sparse types, got: ",
115 options);
116 TORCH_CHECK(
117 at::isFloatingType(typeMetaToScalarType(options.dtype())) || at::isComplexType(typeMetaToScalarType(options.dtype())),
118 function_name,
119 " expects floating point dtypes, got: ",
120 options);
121 TORCH_CHECK(
122 window_length >= 0,
123 function_name,
124 " requires non-negative window_length, got window_length=",
125 window_length);
126 }
127
128 } // namespace
129
130 DEFINE_DISPATCH(complex_stub);
131 DEFINE_DISPATCH(polar_stub);
132
133 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ arange ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
134
arange(const Scalar & end,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)135 Tensor arange(const Scalar& end,
136 std::optional<ScalarType> dtype,
137 std::optional<Layout> layout,
138 std::optional<Device> device,
139 std::optional<bool> pin_memory) {
140 return native::arange(/*start=*/0, end, dtype, layout, device, pin_memory);
141 }
142
arange(const Scalar & start,const Scalar & end,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)143 Tensor arange(const Scalar& start, const Scalar& end,
144 std::optional<ScalarType> dtype,
145 std::optional<Layout> layout,
146 std::optional<Device> device,
147 std::optional<bool> pin_memory) {
148 return native::arange(
149 start, end, /*step=*/1, dtype, layout, device, pin_memory);
150 }
151
arange(const Scalar & start,const Scalar & end,const Scalar & step,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)152 Tensor arange(
153 const Scalar& start,
154 const Scalar& end,
155 const Scalar& step,
156 std::optional<ScalarType> dtype,
157 std::optional<Layout> layout,
158 std::optional<Device> device,
159 std::optional<bool> pin_memory) {
160 // See [Note: hacky wrapper removal for TensorOptions]
161 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
162
163 bool set_to_integral_dtype = !options.has_dtype() &&
164 // bool inputs are considered integral
165 start.isIntegral(true) &&
166 end.isIntegral(true) &&
167 step.isIntegral(true);
168
169 Tensor result = set_to_integral_dtype
170 ? at::empty({0}, options.dtype(at::ScalarType::Long))
171 : at::empty({0}, options);
172 return at::arange_out(result, start, end, step);
173 }
174
arange_out(const Scalar & end,Tensor & result)175 Tensor& arange_out(const Scalar& end, Tensor& result) {
176 return at::arange_out(result, /*start=*/0, end, /*step=*/1);
177 }
178
_dim_arange(const Tensor & like,int64_t dim)179 Tensor _dim_arange(const Tensor& like, int64_t dim) {
180 return at::arange(like.size(dim), like.options().dtype(at::kLong));
181 }
182
183 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ complex / polar ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
184
complex_check_floating(const Tensor & a,const Tensor & b)185 static void complex_check_floating(const Tensor& a, const Tensor& b) {
186 TORCH_CHECK((a.scalar_type() == kFloat || a.scalar_type() == kDouble || a.scalar_type() == kHalf) &&
187 (b.scalar_type() == kFloat || b.scalar_type() == kDouble || b.scalar_type() == kHalf),
188 "Expected both inputs to be Half, Float or Double tensors but got ",
189 a.scalar_type(), " and ", b.scalar_type());
190 }
191
complex_check_dtype(const Tensor & result,const Tensor & a,const Tensor & b)192 static void complex_check_dtype(
193 const Tensor& result,
194 const Tensor& a,
195 const Tensor& b) {
196 complex_check_floating(a, b);
197 TORCH_CHECK(a.scalar_type() == b.scalar_type(),
198 "Expected object of scalar type ", a.scalar_type(),
199 " but got scalar type ", b.scalar_type(), " for second argument");
200 TORCH_CHECK(result.scalar_type() == toComplexType(a.scalar_type()),
201 "Expected object of scalar type ", toComplexType(a.scalar_type()),
202 " but got scalar type ", result.scalar_type(),
203 " for argument 'out'");
204 }
205
complex_out(const Tensor & real,const Tensor & imag,Tensor & result)206 Tensor& complex_out(const Tensor& real, const Tensor& imag, Tensor& result) {
207 complex_check_dtype(result, real, imag);
208 auto iter = TensorIteratorConfig()
209 .add_output(result)
210 .add_const_input(real)
211 .add_const_input(imag)
212 .check_all_same_dtype(false)
213 .build();
214 complex_stub(iter.device_type(), iter);
215 return result;
216 }
217
complex(const Tensor & real,const Tensor & imag)218 Tensor complex(const Tensor& real, const Tensor& imag) {
219 complex_check_floating(real, imag);
220 c10::TensorOptions options = real.options();
221 options = options.dtype(toComplexType(real.scalar_type()));
222 Tensor result = at::empty(0, options);
223 return at::complex_out(result, real, imag);
224 }
225
polar_out(const Tensor & abs,const Tensor & angle,Tensor & result)226 Tensor& polar_out(const Tensor& abs, const Tensor& angle, Tensor& result) {
227 complex_check_dtype(result, abs, angle);
228 auto iter = TensorIteratorConfig()
229 .add_output(result)
230 .add_const_input(abs)
231 .add_const_input(angle)
232 .check_all_same_dtype(false)
233 .build();
234 polar_stub(iter.device_type(), iter);
235 return result;
236 }
237
polar(const Tensor & abs,const Tensor & angle)238 Tensor polar(const Tensor& abs, const Tensor& angle) {
239 complex_check_floating(abs, angle);
240 c10::TensorOptions options = abs.options();
241 options = options.dtype(toComplexType(abs.scalar_type()));
242 Tensor result = at::empty(0, options);
243 return at::polar_out(result, abs, angle);
244 }
245
246 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ empty ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
empty_cpu(IntArrayRef size,std::optional<ScalarType> dtype_opt,std::optional<Layout> layout_opt,std::optional<Device> device_opt,std::optional<bool> pin_memory_opt,std::optional<c10::MemoryFormat> memory_format_opt)247 Tensor empty_cpu(IntArrayRef size, std::optional<ScalarType> dtype_opt, std::optional<Layout> layout_opt,
248 std::optional<Device> device_opt, std::optional<bool> pin_memory_opt, std::optional<c10::MemoryFormat> memory_format_opt) {
249 Tensor result = at::detail::empty_cpu(size, dtype_opt, layout_opt, device_opt, pin_memory_opt, memory_format_opt);
250 // See Note [Enabling Deterministic Operations]
251 if (C10_UNLIKELY(at::globalContext().deterministicAlgorithms() && at::globalContext().deterministicFillUninitializedMemory())) {
252 fill_empty_deterministic_(result);
253 }
254 return result;
255 }
256
empty_names(IntArrayRef size,std::optional<DimnameList> names,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory,std::optional<MemoryFormat> optional_memory_format)257 Tensor empty_names(
258 IntArrayRef size,
259 std::optional<DimnameList> names,
260 std::optional<ScalarType> dtype,
261 std::optional<Layout> layout,
262 std::optional<Device> device,
263 std::optional<bool> pin_memory,
264 std::optional<MemoryFormat> optional_memory_format) {
265 // See [Note: hacky wrapper removal for TensorOptions]
266 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
267
268 if (!names.has_value()) {
269 return at::empty(size, options, optional_memory_format);
270 }
271 TORCH_CHECK(options.layout() == Layout::Strided,
272 "NYI: named tensors only support strided layout");
273 TORCH_CHECK(options.device().is_cpu() || options.device().is_cuda() || options.device().is_xpu() || options.device().is_privateuseone(),
274 "NYI: named tensors only support CPU, CUDA, XPU or ", c10::get_privateuse1_backend(), " tensors.");
275 auto result = at::empty(size, options, optional_memory_format);
276 internal_set_names_inplace(result, names);
277 return result;
278 }
279
empty_permuted_symint(SymIntArrayRef size,IntArrayRef physical_layout,std::optional<ScalarType> dtype_opt,std::optional<Layout> layout_opt,std::optional<Device> device_opt,std::optional<bool> pin_memory_opt)280 Tensor empty_permuted_symint(SymIntArrayRef size, IntArrayRef physical_layout, std::optional<ScalarType> dtype_opt,
281 std::optional<Layout> layout_opt, std::optional<Device> device_opt, std::optional<bool> pin_memory_opt
282 ) {
283 // size is logical; aka, the output size you'll get from the operation overall
284 //
285 // physical_layout follows NCHW/NHWC convention:
286 // contiguous is [0,1,2,3], channels last is [0,2,3,1]
287 //
288 // this means if i is physical index, physical_layout[i] is logical index;
289 // e.g., to find what is innermost physical dim (3), query NHWC[3] == 1
290 // (aka it is channels)
291 int64_t dim = static_cast<int64_t>(size.size());
292 SymDimVector phys_size(dim);
293 TORCH_CHECK(static_cast<int64_t>(physical_layout.size()) == dim,
294 "Number of dimensions in size does not match the "
295 "length of the physical_layout; i.e. len(size) = ", dim,
296 " is not equal to len(physical_layout) = ", physical_layout.size());
297 std::vector<bool> seen_dims(dim);
298 for (const auto i : c10::irange(dim)) {
299 TORCH_CHECK(physical_layout[i] >= 0 && physical_layout[i] < dim,
300 "Dimension out of range (expected to be between 0 and ", dim - 1, ", but got ",
301 physical_layout[i], " at index ", i, "). NB: negative dims "
302 "not currently supported; file an issue if you want it.");
303 TORCH_CHECK(!seen_dims[physical_layout[i]], "Duplicate dim not allowed");
304 phys_size[i] = size[physical_layout[i]];
305 seen_dims[physical_layout[i]] = true;
306 }
307 // do a contiguous allocation
308 Tensor phys_tensor = at::empty_symint(phys_size, dtype_opt, layout_opt, device_opt, pin_memory_opt, std::nullopt);
309 SymIntArrayRef phys_strides = phys_tensor.sym_strides();
310 // permute the strides (inverse permutation! This is why this is
311 // empty_permute*d*, not empty_permute; it's not an empty + permute)
312 SymDimVector strides(dim);
313 for (const auto i : c10::irange(dim)) {
314 strides[physical_layout[i]] = phys_strides[i];
315 }
316 return phys_tensor.as_strided_symint(size, strides);
317 }
318
empty_strided_cpu(IntArrayRef size,IntArrayRef stride,std::optional<ScalarType> dtype_opt,std::optional<Layout> layout_opt,std::optional<Device> device_opt,std::optional<bool> pin_memory_opt)319 Tensor empty_strided_cpu(IntArrayRef size, IntArrayRef stride, std::optional<ScalarType> dtype_opt,
320 std::optional<Layout> layout_opt, std::optional<Device> device_opt, std::optional<bool> pin_memory_opt) {
321 Tensor result = at::detail::empty_strided_cpu(size, stride, dtype_opt, layout_opt, device_opt, pin_memory_opt);
322 // See Note [Enabling Deterministic Operations]
323 if (C10_UNLIKELY(at::globalContext().deterministicAlgorithms() && at::globalContext().deterministicFillUninitializedMemory())) {
324 fill_empty_deterministic_(result);
325 }
326 return result;
327 }
328
empty_out(IntArrayRef size,std::optional<c10::MemoryFormat> optional_memory_format,Tensor & result)329 Tensor& empty_out(IntArrayRef size,
330 std::optional<c10::MemoryFormat> optional_memory_format,
331 Tensor& result) {
332 // Preferably, this argument would not be accepted by _out, but the code
333 // generator requires the out and non-out overloads to match exactly
334 TORCH_CHECK(
335 !optional_memory_format.has_value(),
336 "'memory_format' argument is incompatible with 'out' tensor argument");
337 check_size_nonnegative(size);
338 if (result.is_sparse()) {
339 result.sparse_resize_and_clear_(size, size.size(), 0);
340 } else {
341 result.resize_(size);
342 }
343 // See Note [Enabling Deterministic Operations]
344 if (C10_UNLIKELY(at::globalContext().deterministicAlgorithms() && at::globalContext().deterministicFillUninitializedMemory())) {
345 fill_empty_deterministic_(result);
346 }
347 return result;
348 }
349
350 // Temporary type cast operators. These are needed to trace type-casts now since
351 // Type's are not supported in the IR. Instead, we call down to these
352 // specialized operators for each datatype.
353 // TODO: remove when we have Type support in the IR
354
355 #define DEFINE_CAST_OP(_1, n) \
356 Tensor _cast_##n(const Tensor& self, bool non_blocking) { \
357 if (self.scalar_type() == ScalarType::n) \
358 return self; \
359 return self.to(ScalarType::n, non_blocking); \
360 }
361
362 // Some scalar types in CAST_OP have no declarations, they may be unused in Pytorch.
363 // But we keep them and ignore the warning here until verified in the future.
364 C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wmissing-prototypes")
AT_FORALL_SCALAR_TYPES_AND3(Bool,Half,BFloat16,DEFINE_CAST_OP)365 AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, DEFINE_CAST_OP)
366 C10_DIAGNOSTIC_POP()
367
368 #undef DEFINE_CAST_OP
369
370 Tensor empty_like(
371 const Tensor& self,
372 std::optional<ScalarType> dtype,
373 std::optional<Layout> layout,
374 std::optional<Device> device,
375 std::optional<bool> pin_memory,
376 std::optional<c10::MemoryFormat> optional_memory_format) {
377 // See [Note: hacky wrapper removal for TensorOptions]
378 TensorOptions options_ = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
379
380 TensorOptions options =
381 self.options()
382 .merge_in(options_)
383 .merge_memory_format(optional_memory_format);
384
385 TORCH_CHECK(
386 !(options.layout() != kStrided &&
387 optional_memory_format.has_value()),
388 "memory format option is only supported by strided tensors");
389
390 auto memory_format = options.memory_format_opt().value_or(MemoryFormat::Preserve);
391
392 Tensor result;
393
394 if (memory_format == MemoryFormat::Preserve) {
395 if (self.is_non_overlapping_and_dense()) {
396 result = at::empty_strided_symint(self.sym_sizes(), self.sym_strides(), options.memory_format(std::nullopt));
397 } else if (self.unsafeGetTensorImpl()->support_as_strided() && self.layout() == kStrided) {
398 // If input tensor is not dense and non-overlapping but strided, we will infer an output strides
399 // which keeps the layout permutation of the input tensor.
400 std::vector<int64_t> strides = infer_dense_strides(self.sizes(), self.strides());
401 // See Note [Explicit nullopt MemoryFormat argument]
402 result = at::empty_strided(self.sizes(), strides, options.memory_format(std::nullopt));
403 } else {
404 // See Note [Explicit nullopt MemoryFormat argument]
405 result = at::empty_symint(self.sym_sizes(), options.memory_format(self.suggest_memory_format()), std::nullopt);
406 }
407 } else {
408 // See Note [Explicit nullopt MemoryFormat argument]
409 result = at::empty_symint(self.sym_sizes(), options.memory_format(memory_format), std::nullopt);
410 }
411
412 if (self.opt_names()) {
413 namedinference::propagate_names(result, self.names());
414 }
415
416 // never propagate Conjugate, Negative, and ZeroTensor dispatch key
417 result._set_conj(false);
418 result._set_neg(false);
419 result._set_zero(false);
420 return result;
421 }
422
empty_like_quantized(const Tensor & self,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory,std::optional<c10::MemoryFormat> optional_memory_format)423 Tensor empty_like_quantized(
424 const Tensor& self,
425 std::optional<ScalarType> dtype,
426 std::optional<Layout> layout,
427 std::optional<Device> device,
428 std::optional<bool> pin_memory,
429 std::optional<c10::MemoryFormat> optional_memory_format) {
430 // See [Note: hacky wrapper removal for TensorOptions]
431 TensorOptions options_ = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
432
433 TORCH_CHECK(
434 !(options_.has_memory_format() && optional_memory_format.has_value()),
435 "Cannot set memory_format both in TensorOptions and explicit argument; please delete "
436 "the redundant setter.");
437
438 TensorOptions options =
439 self.options()
440 .merge_in(options_)
441 .merge_memory_format(optional_memory_format);
442
443 TORCH_CHECK(
444 !(options.layout() != kStrided &&
445 optional_memory_format.has_value()),
446 "memory format option is only supported by strided tensors");
447
448 auto memory_format = options.memory_format_opt().value_or(MemoryFormat::Preserve);
449
450
451 // TODO: To support all features of MemoryFormat::Preserve we need to add
452 // _empty_affine_quantized_strided function and use it similarly to
453 // Tensor clone(const Tensor& src, std::optional<c10::MemoryFormat> optional_memory_format)
454 // if (self.is_non_overlapping_and_dense()) -> _empty_affine_quantized_strided
455 if (memory_format == MemoryFormat::Preserve) {
456 memory_format = self.suggest_memory_format();
457 }
458
459
460 // Note [Explicit nullopt MemoryFormat argument]
461 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
462 // Some functions which we call default the OPTIONAL MemoryFormat
463 // argument to something that's not nullopt. If we pass the
464 // MemoryFormat via TensorOptions, we must explicitly disable this
465 // defaulting process, by explicitly passing nullopt for the MemoryFormat
466 // argument. When codegen is adjusted so we can delete this argument from
467 // the method signature, the argument will just disappear entirely.
468 //
469 // BTW, there are a few places where the optional MemoryFormat is None,
470 // but I still pass in nullopt for robustness.
471
472 // We could check if dtype is still quantized? But then should we shift/scale
473 // the q_zero_point / q_scale or not?
474 TORCH_CHECK(!options.has_dtype() || options.dtype() == self.dtype(),
475 "It is currently not supported to specify a dtype that doesn't match "
476 "the input tensor's dtype via empty_like. Specified: ", options.dtype(),
477 " Input tensor's dtype: ", self.dtype());
478 auto qscheme = self.qscheme();
479 if (qscheme == kPerTensorAffine) {
480 return at::_empty_affine_quantized(self.sizes(), options.memory_format(memory_format),
481 self.q_scale(),
482 self.q_zero_point(),
483 // See Note [Explicit nullopt MemoryFormat argument]
484 std::nullopt);
485 } else if (qscheme == kPerChannelAffine) {
486 // Copy the tensors with channels to avoid accidental overrides
487 return at::_empty_per_channel_affine_quantized(
488 self.sizes(),
489 self.q_per_channel_scales().clone(at::MemoryFormat::Preserve),
490 self.q_per_channel_zero_points().clone(at::MemoryFormat::Preserve),
491 self.q_per_channel_axis(),
492 options.memory_format(memory_format),
493 // See Note [Explicit nullopt MemoryFormat argument]
494 std::nullopt);
495 } else {
496 TORCH_CHECK(false, "Unsupported qscheme: ", toString(qscheme));
497 }
498 }
499
new_empty_symint(const Tensor & self,SymIntArrayRef size,std::optional<ScalarType> dtype_opt,std::optional<Layout> layout_opt,std::optional<Device> device_opt,std::optional<bool> pin_memory_opt)500 Tensor new_empty_symint(
501 const Tensor& self,
502 SymIntArrayRef size,
503 std::optional<ScalarType> dtype_opt,
504 std::optional<Layout> layout_opt,
505 std::optional<Device> device_opt,
506 std::optional<bool> pin_memory_opt
507 ) {
508 auto dtype = dtype_opt.has_value() ? dtype_opt : optTypeMetaToScalarType(self.options().dtype_opt());
509 auto layout = layout_opt.has_value() ? layout_opt : self.options().layout_opt();
510 auto device = device_opt.has_value() ? device_opt : self.options().device_opt();
511 auto pin_memory = pin_memory_opt.has_value() ? pin_memory_opt : self.options().pinned_memory_opt();
512 return at::empty_symint(size, dtype, layout, device, pin_memory, std::nullopt);
513 }
514
new_empty_strided_symint(const Tensor & self,c10::SymIntArrayRef size,c10::SymIntArrayRef stride,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)515 Tensor new_empty_strided_symint(
516 const Tensor& self,
517 c10::SymIntArrayRef size,
518 c10::SymIntArrayRef stride,
519 std::optional<ScalarType> dtype,
520 std::optional<Layout> layout,
521 std::optional<Device> device,
522 std::optional<bool> pin_memory
523 ) {
524 // See [Note: hacky wrapper removal for TensorOptions]
525 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
526
527 return at::empty_strided_symint(size, stride, self.options().merge_in(options));
528 }
529
530 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ eye ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
531
eye(int64_t n,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)532 Tensor eye(int64_t n,
533 std::optional<ScalarType> dtype,
534 std::optional<Layout> layout,
535 std::optional<Device> device,
536 std::optional<bool> pin_memory) {
537 // the default value of `m` equals to `n`
538 return at::eye(n, n, dtype, layout, device, pin_memory);
539 }
540
eye(int64_t n,int64_t m,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)541 Tensor eye(int64_t n, int64_t m,
542 std::optional<ScalarType> dtype,
543 std::optional<Layout> layout,
544 std::optional<Device> device,
545 std::optional<bool> pin_memory) {
546 // See [Note: hacky wrapper removal for TensorOptions]
547 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
548
549 auto tensor = at::empty({0}, options); // to be resized
550 return at::eye_out(tensor, n, m);
551 }
552
eye_out_cpu(int64_t n,Tensor & result)553 Tensor& eye_out_cpu(int64_t n, Tensor& result) {
554 // the default value of `m` equals to `n`
555 return native::eye_out_cpu(n, n, result);
556 }
557
eye_out_cpu(int64_t n,int64_t m,Tensor & result)558 Tensor& eye_out_cpu(int64_t n, int64_t m, Tensor& result) {
559 TORCH_CHECK(n >= 0, "n must be greater or equal to 0, got ", n);
560 TORCH_CHECK(m >= 0, "m must be greater or equal to 0, got ", m);
561
562 result.resize_({n, m});
563
564 if (result.is_meta()) return result;
565
566 result.zero_();
567
568 int64_t sz = std::min<int64_t>(n, m);
569 AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND3(kBFloat16, kHalf, kBool, result.scalar_type(), "eye", [&]() -> void {
570 scalar_t* result_data = result.data_ptr<scalar_t>();
571 at::parallel_for(0, sz, internal::GRAIN_SIZE, [&](int64_t p_begin, int64_t p_end) {
572 for (const auto i : c10::irange(p_begin, p_end))result_data[i*(result.strides()[0] + result.strides()[1])] = 1;
573 });
574 });
575
576 return result;
577 }
578
579 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ full ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
580
581 namespace {
582
583 // Performs dtype inference for full
infer_full_options(const Scalar & fill_value,const TensorOptions & options)584 TensorOptions infer_full_options(
585 const Scalar& fill_value,
586 const TensorOptions& options) {
587
588 if (!options.has_dtype()) {
589 if (fill_value.isBoolean()) {
590 return options.dtype(at::kBool);
591 } else if (fill_value.isIntegral(false)) {
592 return options.dtype(at::kLong);
593 } else if (fill_value.isComplex()) {
594 auto scalar_type = (get_default_dtype() == ScalarType::Double) ?
595 ScalarType::ComplexDouble :
596 ScalarType::ComplexFloat;
597 return options.dtype(scalar_type);
598 } else {
599 return options.dtype(get_default_dtype());
600 }
601 }
602
603 return options;
604 }
605
606 } // anonymous namespace
607
full(IntArrayRef size,const Scalar & fill_value,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)608 Tensor full(IntArrayRef size, const Scalar& fill_value,
609 std::optional<ScalarType> dtype,
610 std::optional<Layout> layout,
611 std::optional<Device> device,
612 std::optional<bool> pin_memory) {
613 // See [Note: hacky wrapper removal for TensorOptions]
614 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
615
616 TORCH_CHECK(options.layout() != kSparse,
617 "full(...) is not implemented for sparse layout");
618
619 auto result = at::empty(size, infer_full_options(fill_value, options));
620 return result.fill_(fill_value);
621 }
622
full_out(IntArrayRef size,const Scalar & fill_value,Tensor & result)623 Tensor& full_out(IntArrayRef size, const Scalar& fill_value, Tensor& result) {
624 TORCH_CHECK(!result.is_sparse(),
625 "full(...) is not implemented for sparse layout");
626
627 result.resize_(size);
628 return result.fill_(fill_value);
629 }
630
full_like(const Tensor & self,const Scalar & fill_value,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory,std::optional<c10::MemoryFormat> optional_memory_format)631 Tensor full_like(
632 const Tensor& self,
633 const Scalar& fill_value,
634 std::optional<ScalarType> dtype,
635 std::optional<Layout> layout,
636 std::optional<Device> device,
637 std::optional<bool> pin_memory,
638 std::optional<c10::MemoryFormat> optional_memory_format) {
639 // See [Note: hacky wrapper removal for TensorOptions]
640 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
641
642 auto result = at::empty_like(self, options, optional_memory_format);
643 return result.fill_(fill_value);
644 }
645
new_full(const Tensor & self,IntArrayRef size,const Scalar & fill_value,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)646 Tensor new_full(
647 const Tensor& self,
648 IntArrayRef size,
649 const Scalar& fill_value,
650 std::optional<ScalarType> dtype,
651 std::optional<Layout> layout,
652 std::optional<Device> device,
653 std::optional<bool> pin_memory
654 ) {
655
656 Tensor r = self.new_empty(size, TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory));
657 r.fill_(fill_value);
658 return r;
659 }
660
661 namespace {
linspace_logspace_infer_options(const Scalar & start,const Scalar & end,const TensorOptions & options,const char * fn_name)662 TensorOptions linspace_logspace_infer_options(
663 const Scalar& start,
664 const Scalar& end,
665 const TensorOptions& options,
666 const char* fn_name) {
667 if (start.isComplex() || end.isComplex()) {
668 const auto default_complex_dtype = c10::get_default_complex_dtype();
669 if (options.has_dtype()) {
670 auto dtype = c10::typeMetaToScalarType(options.dtype());
671 TORCH_CHECK(at::isComplexType(dtype),
672 fn_name, ": inferred dtype ", default_complex_dtype, " can't be safely cast to passed dtype ", dtype);
673 } else {
674 return options.dtype(default_complex_dtype);
675 }
676 }
677
678 return options.has_dtype() ? options : options.dtype(c10::get_default_dtype());
679 }
680 } // anonymous namespace
681
682 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ linspace ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
683
linspace(const Scalar & start,const Scalar & end,int64_t steps,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)684 Tensor linspace(
685 const Scalar& start,
686 const Scalar& end,
687 int64_t steps,
688 std::optional<ScalarType> dtype,
689 std::optional<Layout> layout,
690 std::optional<Device> device,
691 std::optional<bool> pin_memory) {
692 // See [Note: hacky wrapper removal for TensorOptions]
693 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
694
695 TORCH_CHECK(steps >= 0, "number of steps must be non-negative");
696 auto result_options = linspace_logspace_infer_options(start, end, options, "torch.linspace()");
697 Tensor result = at::empty({steps}, result_options);
698 return at::linspace_out(result, start, end, steps);
699 }
700
linspace(const Tensor & start,const Tensor & end,int64_t steps,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)701 Tensor linspace(
702 const Tensor& start,
703 const Tensor& end,
704 int64_t steps,
705 std::optional<ScalarType> dtype,
706 std::optional<Layout> layout,
707 std::optional<Device> device,
708 std::optional<bool> pin_memory) {
709 TORCH_CHECK(start.dim() == 0 && end.dim() == 0, "linspace only supports 0-dimensional start and end tensors, "
710 "but got start with ", start.dim(), " dimension(s) and end with ", end.dim()," dimension(s).");
711 return at::linspace(start.item(), end.item(), steps, dtype, layout, device, pin_memory);
712 }
713
linspace(const Tensor & start,const Scalar & end,int64_t steps,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)714 Tensor linspace(
715 const Tensor& start,
716 const Scalar& end,
717 int64_t steps,
718 std::optional<ScalarType> dtype,
719 std::optional<Layout> layout,
720 std::optional<Device> device,
721 std::optional<bool> pin_memory) {
722 TORCH_CHECK(start.dim() == 0, "linspace only supports 0-dimensional start and end tensors, "
723 "but got start with ", start.dim(), " dimension(s).");
724 return at::linspace(start.item(), end, steps, dtype, layout, device, pin_memory);
725 }
726
linspace(const Scalar & start,const Tensor & end,int64_t steps,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)727 Tensor linspace(
728 const Scalar& start,
729 const Tensor& end,
730 int64_t steps,
731 std::optional<ScalarType> dtype,
732 std::optional<Layout> layout,
733 std::optional<Device> device,
734 std::optional<bool> pin_memory) {
735 TORCH_CHECK(end.dim() == 0, "linspace only supports 0-dimensional start and end tensors, "
736 "but got end with ", end.dim()," dimension(s).");
737 return at::linspace(start, end.item(), steps, dtype, layout, device, pin_memory);
738 }
739
740 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ logspace ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
741
logspace(const Scalar & start,const Scalar & end,int64_t steps,double base,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)742 Tensor logspace(
743 const Scalar& start,
744 const Scalar& end,
745 int64_t steps,
746 double base,
747 std::optional<ScalarType> dtype,
748 std::optional<Layout> layout,
749 std::optional<Device> device,
750 std::optional<bool> pin_memory) {
751 // See [Note: hacky wrapper removal for TensorOptions]
752 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
753
754 TORCH_CHECK(steps >= 0, "number of steps must be non-negative");
755 auto result_options = linspace_logspace_infer_options(start, end, options, "torch.logspace()");
756 Tensor result = at::empty({steps}, result_options);
757 return at::logspace_out(result, start, end, steps, base);
758 }
759
logspace(const Tensor & start,const Tensor & end,int64_t steps,double base,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)760 Tensor logspace(
761 const Tensor& start,
762 const Tensor& end,
763 int64_t steps,
764 double base,
765 std::optional<ScalarType> dtype,
766 std::optional<Layout> layout,
767 std::optional<Device> device,
768 std::optional<bool> pin_memory) {
769 TORCH_CHECK(start.dim() == 0 && end.dim() == 0, "logspace only supports 0-dimensional start and end tensors, "
770 "but got start with ", start.dim(), " dimension(s) and end with ", end.dim()," dimension(s).");
771 return at::logspace(start.item(), end.item(), steps, base, dtype, layout, device, pin_memory);
772 }
773
logspace(const Tensor & start,const Scalar & end,int64_t steps,double base,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)774 Tensor logspace(
775 const Tensor& start,
776 const Scalar& end,
777 int64_t steps,
778 double base,
779 std::optional<ScalarType> dtype,
780 std::optional<Layout> layout,
781 std::optional<Device> device,
782 std::optional<bool> pin_memory) {
783 TORCH_CHECK(start.dim() == 0, "logspace only supports 0-dimensional start and end tensors, "
784 "but got start with ", start.dim(), " dimension(s).");
785 return at::logspace(start.item(), end, steps, base, dtype, layout, device, pin_memory);
786 }
787
logspace(const Scalar & start,const Tensor & end,int64_t steps,double base,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)788 Tensor logspace(
789 const Scalar& start,
790 const Tensor& end,
791 int64_t steps,
792 double base,
793 std::optional<ScalarType> dtype,
794 std::optional<Layout> layout,
795 std::optional<Device> device,
796 std::optional<bool> pin_memory) {
797 TORCH_CHECK(end.dim() == 0, "logspace only supports 0-dimensional start and end tensors, "
798 "but got end with ", end.dim()," dimension(s).");
799 return at::logspace(start, end.item(), steps, base, dtype, layout, device, pin_memory);
800 }
801
802 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ones ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
803
ones(IntArrayRef size,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)804 Tensor ones(IntArrayRef size,
805 std::optional<ScalarType> dtype,
806 std::optional<Layout> layout,
807 std::optional<Device> device,
808 std::optional<bool> pin_memory) {
809 return native::full(size, /*fill_value=*/1., dtype, layout, device, pin_memory);
810 }
811
ones_out(IntArrayRef size,Tensor & result)812 Tensor& ones_out(IntArrayRef size, Tensor& result) {
813 return native::full_out(size, /*fill_value=*/1., result);
814 }
815
ones_like(const Tensor & self,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory,std::optional<c10::MemoryFormat> optional_memory_format)816 Tensor ones_like(
817 const Tensor& self,
818 std::optional<ScalarType> dtype,
819 std::optional<Layout> layout,
820 std::optional<Device> device,
821 std::optional<bool> pin_memory,
822 std::optional<c10::MemoryFormat> optional_memory_format) {
823 auto result = at::empty_like(self, dtype, layout, device, pin_memory, optional_memory_format);
824 return result.fill_(1.);
825 }
826
new_ones(const Tensor & self,IntArrayRef size,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)827 Tensor new_ones(
828 const Tensor& self,
829 IntArrayRef size,
830 std::optional<ScalarType> dtype,
831 std::optional<Layout> layout,
832 std::optional<Device> device,
833 std::optional<bool> pin_memory) {
834 // See [Note: hacky wrapper removal for TensorOptions]
835 Tensor r = self.new_empty(size, TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory));
836 r.fill_(1.);
837 return r;
838 }
839
840 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ scalar_tensor ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
841
scalar_tensor(const Scalar & s,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)842 Tensor scalar_tensor(const Scalar& s,
843 std::optional<ScalarType> dtype,
844 std::optional<Layout> layout,
845 std::optional<Device> device,
846 std::optional<bool> pin_memory) {
847 // See [Note: hacky wrapper removal for TensorOptions]
848 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
849
850 if (options.device() == at::kCPU) {
851 // This is a fast track to skip device dispatch for making scalar tensor on CPU.
852 // See https://github.com/pytorch/pytorch/pull/29915 for more detailed perf
853 // difference.
854 // In the future when we remove the overhead of device dispatch, we'll happily
855 // revert this to following:
856 // auto result = at::empty({}, options);
857 at::tracer::impl::NoTracerDispatchMode tracer_guard;
858 at::AutoDispatchBelowAutograd mode;
859 auto result = empty_cpu({}, optTypeMetaToScalarType(options.dtype_opt()), options.layout_opt(), options.device_opt(), options.pinned_memory_opt());
860 at::native::fill_(result, s);
861 return result;
862 }
863 return at::empty({}, options).fill_(s);
864 }
865
866 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ rand ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
867
rand(IntArrayRef size,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)868 Tensor rand(IntArrayRef size,
869 std::optional<ScalarType> dtype,
870 std::optional<Layout> layout,
871 std::optional<Device> device,
872 std::optional<bool> pin_memory) {
873 return native::rand(size, static_cast<std::optional<Generator>>(std::nullopt), dtype, layout, device, pin_memory);
874 }
875
rand(IntArrayRef size,std::optional<Generator> generator,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)876 Tensor rand(IntArrayRef size, std::optional<Generator> generator,
877 std::optional<ScalarType> dtype,
878 std::optional<Layout> layout,
879 std::optional<Device> device,
880 std::optional<bool> pin_memory) {
881 // See [Note: hacky wrapper removal for TensorOptions]
882 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
883
884 auto result = at::empty(size, options);
885 return result.uniform_(0, 1, std::move(generator));
886 }
887
rand_out(IntArrayRef size,Tensor & result)888 Tensor& rand_out(IntArrayRef size, Tensor& result) {
889 return native::rand_out(size, std::nullopt, result);
890 }
891
rand_out(IntArrayRef size,std::optional<Generator> generator,Tensor & result)892 Tensor& rand_out(IntArrayRef size, std::optional<Generator> generator, Tensor& result) {
893 result.resize_(size);
894 return result.uniform_(0, 1, std::move(generator));
895 }
896
rand_like(const Tensor & self,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory,std::optional<c10::MemoryFormat> optional_memory_format)897 Tensor rand_like(
898 const Tensor& self,
899 std::optional<ScalarType> dtype,
900 std::optional<Layout> layout,
901 std::optional<Device> device,
902 std::optional<bool> pin_memory,
903 std::optional<c10::MemoryFormat> optional_memory_format) {
904 // See [Note: hacky wrapper removal for TensorOptions]
905 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
906
907 auto result = at::empty_like(self, options, optional_memory_format);
908 return result.uniform_(0, 1, std::nullopt);
909 }
910
911 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ randint ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
912
randint(int64_t high,IntArrayRef size,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)913 Tensor randint(int64_t high, IntArrayRef size,
914 std::optional<ScalarType> dtype,
915 std::optional<Layout> layout,
916 std::optional<Device> device,
917 std::optional<bool> pin_memory) {
918 return native::randint(high, size, std::nullopt /* generator*/, dtype, layout, device, pin_memory);
919 }
920
randint(int64_t high,IntArrayRef size,std::optional<Generator> generator,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)921 Tensor randint(
922 int64_t high,
923 IntArrayRef size,
924 std::optional<Generator> generator,
925 std::optional<ScalarType> dtype,
926 std::optional<Layout> layout,
927 std::optional<Device> device,
928 std::optional<bool> pin_memory) {
929 return native::randint(0, high, size, std::move(generator), dtype, layout, device, pin_memory);
930 }
931
randint(int64_t low,int64_t high,IntArrayRef size,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)932 Tensor randint(
933 int64_t low,
934 int64_t high,
935 IntArrayRef size,
936 std::optional<ScalarType> dtype,
937 std::optional<Layout> layout,
938 std::optional<Device> device,
939 std::optional<bool> pin_memory) {
940 return native::randint(low, high, size, std::nullopt, dtype, layout, device, pin_memory);
941 }
942
randint(int64_t low,int64_t high,IntArrayRef size,std::optional<Generator> generator,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)943 Tensor randint(
944 int64_t low,
945 int64_t high,
946 IntArrayRef size,
947 std::optional<Generator> generator,
948 std::optional<ScalarType> dtype,
949 std::optional<Layout> layout,
950 std::optional<Device> device,
951 std::optional<bool> pin_memory) {
952 // See [Note: hacky wrapper removal for TensorOptions]
953 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
954
955 auto result = at::empty(size, options);
956 return result.random_(low, high, std::move(generator));
957 }
958
randint_out(int64_t high,IntArrayRef size,Tensor & result)959 Tensor& randint_out(int64_t high, IntArrayRef size, Tensor& result) {
960 return native::randint_out(high, size, std::nullopt, result);
961 }
962
randint_out(int64_t high,IntArrayRef size,std::optional<Generator> generator,Tensor & result)963 Tensor& randint_out(int64_t high,
964 IntArrayRef size,
965 std::optional<Generator> generator,
966 Tensor& result) {
967 result.resize_(size);
968 return result.random_(0, high, std::move(generator));
969 }
970
randint_out(int64_t low,int64_t high,IntArrayRef size,Tensor & result)971 Tensor& randint_out(int64_t low, int64_t high, IntArrayRef size, Tensor& result) {
972 return native::randint_out(low, high, size, std::nullopt, result);
973 }
974
randint_out(int64_t low,int64_t high,IntArrayRef size,std::optional<Generator> generator,Tensor & result)975 Tensor& randint_out(int64_t low,
976 int64_t high,
977 IntArrayRef size,
978 std::optional<Generator> generator,
979 Tensor& result) {
980 result.resize_(size);
981 return result.random_(low, high, std::move(generator));
982 }
983
randint_like(const Tensor & self,int64_t high,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory,std::optional<c10::MemoryFormat> optional_memory_format)984 Tensor randint_like(
985 const Tensor& self,
986 int64_t high,
987 std::optional<ScalarType> dtype,
988 std::optional<Layout> layout,
989 std::optional<Device> device,
990 std::optional<bool> pin_memory,
991 std::optional<c10::MemoryFormat> optional_memory_format) {
992 // See [Note: hacky wrapper removal for TensorOptions]
993 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
994
995 auto result = at::empty_like(self, options, optional_memory_format);
996 return result.random_(0, high, std::nullopt);
997 }
998
randint_like(const Tensor & self,int64_t low,int64_t high,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory,std::optional<c10::MemoryFormat> optional_memory_format)999 Tensor randint_like(
1000 const Tensor& self,
1001 int64_t low,
1002 int64_t high,
1003 std::optional<ScalarType> dtype,
1004 std::optional<Layout> layout,
1005 std::optional<Device> device,
1006 std::optional<bool> pin_memory,
1007 std::optional<c10::MemoryFormat> optional_memory_format) {
1008 // See [Note: hacky wrapper removal for TensorOptions]
1009 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
1010
1011 auto result = at::empty_like(self, options, optional_memory_format);
1012 return result.random_(low, high, std::nullopt);
1013 }
1014
1015 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ randn ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1016
randn(IntArrayRef size,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1017 Tensor randn(IntArrayRef size,
1018 std::optional<ScalarType> dtype,
1019 std::optional<Layout> layout,
1020 std::optional<Device> device,
1021 std::optional<bool> pin_memory) {
1022 return native::randn(size, static_cast<std::optional<Generator>>(std::nullopt), dtype, layout, device, pin_memory);
1023 }
1024
randn(IntArrayRef size,std::optional<Generator> generator,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1025 Tensor randn(IntArrayRef size, std::optional<Generator> generator,
1026 std::optional<ScalarType> dtype,
1027 std::optional<Layout> layout,
1028 std::optional<Device> device,
1029 std::optional<bool> pin_memory) {
1030 // See [Note: hacky wrapper removal for TensorOptions]
1031 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
1032
1033 auto result = at::empty(size, options);
1034 return result.normal_(0, 1, std::move(generator));
1035 }
1036
randn_out(IntArrayRef size,Tensor & result)1037 Tensor& randn_out(IntArrayRef size, Tensor& result) {
1038 return native::randn_out(size, std::nullopt, result);
1039 }
1040
randn_out(IntArrayRef size,std::optional<Generator> generator,Tensor & result)1041 Tensor& randn_out(IntArrayRef size, std::optional<Generator> generator, Tensor& result) {
1042 result.resize_(size);
1043 return result.normal_(0, 1, std::move(generator));
1044 }
1045
normal(double mean,double std,IntArrayRef size,std::optional<Generator> generator,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1046 Tensor normal(double mean, double std, IntArrayRef size,
1047 std::optional<Generator> generator,
1048 std::optional<ScalarType> dtype,
1049 std::optional<Layout> layout,
1050 std::optional<Device> device,
1051 std::optional<bool> pin_memory) {
1052 // See [Note: hacky wrapper removal for TensorOptions]
1053 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
1054
1055 auto result = at::empty(size, options);
1056 return result.normal_(mean, std, std::move(generator));
1057 }
1058
normal_out(double mean,double std,IntArrayRef size,std::optional<Generator> generator,Tensor & result)1059 Tensor& normal_out(double mean, double std,
1060 IntArrayRef size, std::optional<Generator> generator, Tensor& result) {
1061 result.resize_(size);
1062 return result.normal_(mean, std, std::move(generator));
1063 }
1064
randn_like(const Tensor & self,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory,std::optional<c10::MemoryFormat> optional_memory_format)1065 Tensor randn_like(
1066 const Tensor& self,
1067 std::optional<ScalarType> dtype,
1068 std::optional<Layout> layout,
1069 std::optional<Device> device,
1070 std::optional<bool> pin_memory,
1071 std::optional<c10::MemoryFormat> optional_memory_format) {
1072 // See [Note: hacky wrapper removal for TensorOptions]
1073 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
1074
1075 auto result = at::empty_like(self, options, optional_memory_format);
1076 return result.normal_(0, 1, std::nullopt);
1077 }
1078
1079 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ randperm ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1080
1081 namespace {
1082 template <typename scalar_t>
randperm_cpu(Tensor & result,int64_t n,CPUGeneratorImpl * generator)1083 void randperm_cpu(Tensor& result, int64_t n, CPUGeneratorImpl* generator) {
1084 scalar_t *r__data = result.data_ptr<scalar_t>();
1085
1086 result.resize_({n});
1087 int64_t r__stride_0 = result.stride(0);
1088
1089 at::parallel_for(0, n, internal::GRAIN_SIZE,
1090 [&r__data, &r__stride_0](int64_t p_begin, int64_t p_end) {
1091 for (const auto i : c10::irange(p_begin, p_end)) {
1092 r__data[i*r__stride_0] = static_cast<scalar_t>(i);
1093 }
1094 });
1095
1096 for(int64_t i = 0; i < n - 1; i++)
1097 {
1098 // NOLINTNEXTLINE(clang-analyzer-security.insecureAPI.rand)
1099 int64_t z = generator->random() % (n-i);
1100 scalar_t sav = r__data[i*r__stride_0];
1101 r__data[i*r__stride_0] = r__data[(z+i)*r__stride_0];
1102 r__data[(z+i)*r__stride_0] = sav;
1103 }
1104 }
1105 } // namespace
1106
randperm(int64_t n,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1107 Tensor randperm(int64_t n,
1108 std::optional<ScalarType> dtype,
1109 std::optional<Layout> layout,
1110 std::optional<Device> device,
1111 std::optional<bool> pin_memory) {
1112 return native::randperm(n, std::nullopt, dtype, layout, device, pin_memory);
1113 }
1114
randperm(int64_t n,std::optional<Generator> generator,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1115 Tensor randperm(int64_t n, std::optional<Generator> generator,
1116 std::optional<ScalarType> dtype,
1117 std::optional<Layout> layout,
1118 std::optional<Device> device,
1119 std::optional<bool> pin_memory) {
1120 if (!dtype.has_value()) {
1121 dtype = ScalarType::Long;
1122 }
1123
1124 // See [Note: hacky wrapper removal for TensorOptions]
1125 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
1126
1127 auto tensor = at::empty(n, options);
1128 return at::randperm_out(tensor, n, std::move(generator));
1129 }
1130
randperm_out(int64_t n,Tensor & result)1131 Tensor& randperm_out(int64_t n, Tensor& result) {
1132 return at::randperm_out(result, n, std::nullopt);
1133 }
1134
randperm_out_cpu(int64_t n,std::optional<Generator> generator,Tensor & result)1135 Tensor& randperm_out_cpu(int64_t n, std::optional<Generator> generator, Tensor& result) {
1136 TORCH_CHECK(n >= 0, "n must be non-negative, got", n);
1137 TORCH_CHECK(!generator.has_value() || (generator.has_value() && result.device() == generator->device()), "Expected a '", result.device(), "' generator device but found '", generator->device(), "'");
1138 check_supported_max_int_with_precision(n, result);
1139 result.resize_({n});
1140 auto gen = get_generator_or_default<CPUGeneratorImpl>(generator, detail::getDefaultCPUGenerator());
1141 // See Note [Acquire lock when using random generators]
1142 std::lock_guard<std::mutex> lock(gen->mutex_);
1143 AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, result.scalar_type(), "randperm", [&]() -> void {
1144 randperm_cpu<scalar_t>(result, n, gen);
1145 });
1146
1147 return result;
1148 }
1149
1150 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ range ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1151
range(const Scalar & start,const Scalar & end,const Scalar & step,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1152 Tensor range(
1153 const Scalar& start,
1154 const Scalar& end,
1155 const Scalar& step,
1156 std::optional<ScalarType> dtype,
1157 std::optional<Layout> layout,
1158 std::optional<Device> device,
1159 std::optional<bool> pin_memory) {
1160 // See [Note: hacky wrapper removal for TensorOptions]
1161 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
1162
1163 Tensor result = at::empty({0}, options);
1164 return at::range_out(result, start, end, step);
1165 }
1166
range(const Scalar & start,const Scalar & end,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1167 Tensor range(
1168 const Scalar& start,
1169 const Scalar& end,
1170 std::optional<ScalarType> dtype,
1171 std::optional<Layout> layout,
1172 std::optional<Device> device,
1173 std::optional<bool> pin_memory) {
1174 return at::native::range(start, end, 1, dtype, layout, device, pin_memory);
1175 }
1176
1177 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ triangle ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1178
tril_indices_cpu(int64_t row,int64_t col,int64_t offset,std::optional<ScalarType> dtype_opt,std::optional<Layout> layout_opt,std::optional<Device> device_opt,std::optional<bool> pin_memory_opt)1179 Tensor tril_indices_cpu(
1180 int64_t row, int64_t col, int64_t offset, std::optional<ScalarType> dtype_opt,
1181 std::optional<Layout> layout_opt, std::optional<Device> device_opt, std::optional<bool> pin_memory_opt) {
1182 if (!dtype_opt.has_value()) {
1183 dtype_opt = ScalarType::Long;
1184 }
1185
1186 check_args(row, col, layout_opt);
1187
1188 auto tril_size = get_tril_size(row, col, offset);
1189
1190 // create an empty Tensor with correct size
1191 auto result = at::native::empty_cpu({2, tril_size}, dtype_opt, layout_opt, device_opt, pin_memory_opt);
1192
1193 // The following three approaches result in very little performance
1194 // differences. Hence, the 2nd option is taken for simpler code, and to return
1195 // contiguous tensors. Refer to #14904 for more details.
1196 //
1197 // 1. sequential RAM access: fill row coordinates first, then columns. This
1198 // results in two for-loop and more arithmetic operations.
1199 //
1200 // 2. interleaved RAM access: fill in index coordinates one by one, which
1201 // jumps between the two output Tensor rows in every iteration.
1202 //
1203 // 3. sequential RAM + transpose: create an n X 2 Tensor, fill the Tensor
1204 // sequentially, and then transpose it.
1205 AT_DISPATCH_INDEX_TYPES(result.scalar_type(), "tril_indices", [&]() -> void {
1206 // fill the Tensor with correct values
1207 index_t* result_data = result.data_ptr<index_t>();
1208 int64_t i = 0;
1209
1210 index_t r = std::max<int64_t>(0, -offset), c = 0;
1211 while (i < tril_size) {
1212 result_data[i] = r;
1213 result_data[tril_size + i++] = c;
1214
1215 // move to the next column and check if (r, c) is still in bound
1216 c += 1;
1217 if (c > r + offset || c >= col) {
1218 r += 1;
1219 c = 0;
1220 // NOTE: not necessary to check if r is less than row here, because i
1221 // and tril_size provide the guarantee
1222 }
1223 }
1224 });
1225
1226 return result;
1227 }
1228
triu_indices_cpu(int64_t row,int64_t col,int64_t offset,std::optional<ScalarType> dtype_opt,std::optional<Layout> layout_opt,std::optional<Device> device_opt,std::optional<bool> pin_memory_opt)1229 Tensor triu_indices_cpu(
1230 int64_t row, int64_t col, int64_t offset, std::optional<ScalarType> dtype_opt,
1231 std::optional<Layout> layout_opt, std::optional<Device> device_opt, std::optional<bool> pin_memory_opt) {
1232 if (!dtype_opt.has_value()) {
1233 dtype_opt = ScalarType::Long;
1234 }
1235
1236 check_args(row, col, layout_opt);
1237
1238 auto triu_size = row * col - get_tril_size(row, col, offset - 1);
1239
1240 // create an empty Tensor with correct size
1241 auto result = at::native::empty_cpu({2, triu_size}, dtype_opt, layout_opt, device_opt, pin_memory_opt);
1242
1243 AT_DISPATCH_INDEX_TYPES(result.scalar_type(), "triu_indices", [&]() -> void {
1244 // fill the Tensor with correct values
1245 index_t* result_data = result.data_ptr<index_t>();
1246 int64_t i = 0;
1247 // not typing std::max with scalar_t as it could be an unsigned type
1248 // NOTE: no need to check if the returned value of std::max overflows
1249 // index_t, as i and triu_size act as a guard.
1250 index_t c = std::max<int64_t>(0, offset), r = 0;
1251 while (i < triu_size) {
1252 result_data[i] = r;
1253 result_data[triu_size + i++] = c;
1254
1255 // move to the next column and check if (r, c) is still in bound
1256 c += 1;
1257 if (c >= col) {
1258 r += 1;
1259 // not typing std::max with scalar_t as it could be an unsigned type
1260 // NOTE: not necessary to check if c is less than col or overflows here,
1261 // because i and triu_size act as a guard.
1262 c = std::max<int64_t>(0, r + offset);
1263 }
1264 }
1265 });
1266
1267 return result;
1268 }
1269
1270 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ zeros ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1271
zeros_sparse_compressed_symint(c10::SymIntArrayRef size,std::optional<ScalarType> dtype,Layout layout,std::optional<Device> device,std::optional<bool> pin_memory)1272 static Tensor zeros_sparse_compressed_symint(c10::SymIntArrayRef size,
1273 std::optional<ScalarType> dtype,
1274 Layout layout,
1275 std::optional<Device> device,
1276 std::optional<bool> pin_memory) {
1277 check_size_nonnegative(size);
1278 TORCH_CHECK(size.size() >= 2, "torch.zeros: Only batched sparse compressed (non-block) tensors are supported, but got size ", size);
1279 auto size_ = C10_AS_INTARRAYREF_SLOW(size);
1280 // torch.zeros cannot be used to create blocked tensors because its
1281 // API lacks a method to specify the block size.
1282 AT_DISPATCH_SPARSE_COMPRESSED_NONBLOCK_LAYOUTS(layout, "zeros_sparse_compressed", [&]{});
1283
1284 int64_t nnz = 0;
1285 auto compressed_indices_size = DimVector(size_.slice(0, size.size() - 2));
1286 auto plain_indices_and_values_size = DimVector(size_.slice(0, size.size() - 2));
1287 compressed_indices_size.push_back(size_[at::sparse_csr::compressedDimension(layout, size_)] + 1);
1288 plain_indices_and_values_size.push_back(nnz);
1289
1290 TensorOptions options = TensorOptions().dtype(ScalarType::Long).layout(Layout::Strided).device(device).pinned_memory(pin_memory);
1291 auto compressed_indices = at::empty(compressed_indices_size, options);
1292 compressed_indices.zero_();
1293 auto plain_indices = at::empty(plain_indices_and_values_size, options);
1294 auto values = at::empty(plain_indices_and_values_size, options.dtype(dtype));
1295
1296 return at::_sparse_compressed_tensor_unsafe(compressed_indices,
1297 plain_indices,
1298 values,
1299 size_,
1300 dtype,
1301 layout,
1302 device,
1303 pin_memory);
1304 }
1305
zeros_symint(SymIntArrayRef size,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1306 Tensor zeros_symint(SymIntArrayRef size,
1307 std::optional<ScalarType> dtype,
1308 std::optional<Layout> layout,
1309 std::optional<Device> device,
1310 std::optional<bool> pin_memory) {
1311 Layout layout_ = layout.value_or(Layout::Strided);
1312 if (at::sparse_csr::is_sparse_compressed(layout_)) {
1313 return zeros_sparse_compressed_symint(size, dtype, layout_, device, pin_memory);
1314 }
1315 // See [Note: hacky wrapper removal for TensorOptions]
1316 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
1317 auto result = at::empty_symint(size, options);
1318 return result.zero_();
1319 }
1320
_efficientzerotensor(IntArrayRef size,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1321 Tensor _efficientzerotensor(IntArrayRef size,
1322 std::optional<ScalarType> dtype,
1323 std::optional<Layout> layout,
1324 std::optional<Device> device,
1325 std::optional<bool> pin_memory) {
1326 auto device_ = device_or_default(device);
1327 auto allocator = at::native::ZeroTensorAllocator(device_);
1328 auto dtype_ = dtype_or_default(dtype);
1329 auto zero_ks = at::DispatchKeySet(c10::DispatchKey::CPU) | at::DispatchKeySet(c10::DispatchKey::ZeroTensor);
1330 auto out = at::detail::empty_generic(size, &allocator, zero_ks, dtype_, std::nullopt);
1331 return out;
1332 }
1333
_efficientzerotensor_meta_symint(SymIntArrayRef size,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1334 Tensor _efficientzerotensor_meta_symint(SymIntArrayRef size,
1335 std::optional<ScalarType> dtype,
1336 std::optional<Layout> layout,
1337 std::optional<Device> device,
1338 std::optional<bool> pin_memory) {
1339 auto device_ = device_or_default(device);
1340 auto allocator = at::native::ZeroTensorAllocator(device_);
1341 auto dtype_ = dtype_or_default(dtype);
1342 auto zero_ks = at::DispatchKeySet(c10::DispatchKey::Meta) | at::DispatchKeySet(c10::DispatchKey::ZeroTensor);
1343 auto out = at::detail::empty_generic_symint(size, &allocator, zero_ks, dtype_, std::nullopt);
1344 return out;
1345 }
1346
zeros_sparse_out(IntArrayRef size,Tensor & result)1347 Tensor& zeros_sparse_out(IntArrayRef size, Tensor& result) {
1348 result.sparse_resize_and_clear_(size, size.size(), 0.);
1349 return result;
1350 }
1351
zeros_out(IntArrayRef size,Tensor & result)1352 Tensor& zeros_out(IntArrayRef size, Tensor& result) {
1353 if (result.is_sparse()) {
1354 // TODO: I think this branch should be dead, but we don't have an easy
1355 // way to cover all sparse kernels with zeros_sparse_out, so retain this
1356 // for now
1357 result.sparse_resize_and_clear_(size, size.size(), 0.);
1358 return result;
1359 } else {
1360 result.resize_(size);
1361 }
1362 return result.zero_();
1363 }
1364
zeros_like(const Tensor & self,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory,std::optional<c10::MemoryFormat> optional_memory_format)1365 Tensor zeros_like(
1366 const Tensor& self,
1367 std::optional<ScalarType> dtype,
1368 std::optional<Layout> layout,
1369 std::optional<Device> device,
1370 std::optional<bool> pin_memory,
1371 std::optional<c10::MemoryFormat> optional_memory_format) {
1372 // See [Note: hacky wrapper removal for TensorOptions]
1373 auto other_options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
1374 // Prefer values passed in explicitly, but default to value from self.
1375 auto options = self.options().merge_in(other_options);
1376
1377 if (options.layout() == kSparse) {
1378 TORCH_CHECK(
1379 !(optional_memory_format.has_value()),
1380 "memory format option is only supported by strided tensors");
1381 auto res = at::empty({0}, self.options().merge_in(options)); // to be resized
1382
1383 if (self.is_sparse()) {
1384 res.sparse_resize_and_clear_(
1385 self.sizes(), self.sparse_dim(), self.dense_dim());
1386 } else if (at::sparse_csr::is_sparse_compressed(self)) {
1387 res.sparse_resize_and_clear_(
1388 self.sizes(), self.sizes().size() - self.dense_dim(), self.dense_dim());
1389 } else {
1390 res.sparse_resize_and_clear_(self.sizes(), self.sizes().size(), 0);
1391 }
1392 res._coalesced_(true);
1393
1394 return res;
1395 } else if (at::sparse_csr::is_sparse_compressed(options.layout())) {
1396 int64_t nnz = 0;
1397 int64_t dense_dim = (self.layout() == kStrided ? self.dim() - 2: self.dense_dim());
1398 DimVector blocksize{};
1399 if (self.layout() == kSparseBsr || self.layout() == kSparseBsc) {
1400 blocksize.append(at::sparse_csr::getBlockSize(self));
1401 }
1402 ScalarType index_dtype = at::sparse_csr::getIndexDtype(self);
1403 auto res = at::native::sparse_compressed_tensor_with_dims(
1404 nnz, dense_dim, self.sizes(), blocksize, index_dtype,
1405 typeMetaToScalarType(options.dtype()), options.layout(), options.device(), options.pinned_memory());
1406 auto [compressed_indices, plain_indices] = at::sparse_csr::getCompressedPlainIndices(res);
1407 compressed_indices.zero_();
1408 return res;
1409 }
1410 auto result = at::empty_like(self, options, optional_memory_format);
1411 return result.zero_();
1412 }
1413
new_zeros(const Tensor & self,IntArrayRef size,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1414 Tensor new_zeros(
1415 const Tensor& self,
1416 IntArrayRef size,
1417 std::optional<ScalarType> dtype,
1418 std::optional<Layout> layout,
1419 std::optional<Device> device,
1420 std::optional<bool> pin_memory
1421 ) {
1422 Tensor r = self.new_empty(size, TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory));
1423 r.zero_();
1424 return r;
1425 }
1426
1427 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~ bartlett_window ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1428
bartlett_window(int64_t window_length,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1429 Tensor bartlett_window(int64_t window_length,
1430 std::optional<ScalarType> dtype,
1431 std::optional<Layout> layout,
1432 std::optional<Device> device,
1433 std::optional<bool> pin_memory) {
1434 return native::bartlett_window(
1435 window_length, /*periodic=*/true, dtype, layout, device, pin_memory);
1436 }
1437
bartlett_window(int64_t window_length,bool periodic,std::optional<ScalarType> dtype_opt,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1438 Tensor bartlett_window(
1439 int64_t window_length,
1440 bool periodic,
1441 std::optional<ScalarType> dtype_opt,
1442 std::optional<Layout> layout,
1443 std::optional<Device> device,
1444 std::optional<bool> pin_memory) {
1445 // See [Note: hacky wrapper removal for TensorOptions]
1446 ScalarType dtype = c10::dtype_or_default(dtype_opt);
1447 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
1448
1449 window_function_checks("bartlett_window", options, window_length);
1450 if (window_length == 0) {
1451 return at::empty({0}, options);
1452 }
1453 if (window_length == 1) {
1454 return native::ones({1}, dtype, layout, device, pin_memory);
1455 }
1456 if (periodic) {
1457 window_length += 1;
1458 }
1459 auto window = native::arange(window_length, dtype, layout, device, pin_memory)
1460 .mul_(2. / static_cast<double>(window_length - 1));
1461 const int64_t first_half_size = ((window_length - 1) >> 1) + 1;
1462 window.narrow(0, first_half_size, window_length - first_half_size).mul_(-1).add_(2);
1463 return periodic ? window.narrow(0, 0, window_length - 1) : std::move(window);
1464 }
1465
1466 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~ blackman_window ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1467
blackman_window(int64_t window_length,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1468 Tensor blackman_window(int64_t window_length,
1469 std::optional<ScalarType> dtype,
1470 std::optional<Layout> layout,
1471 std::optional<Device> device,
1472 std::optional<bool> pin_memory) {
1473 return native::blackman_window(
1474 window_length, /*periodic=*/true, dtype, layout, device, pin_memory);
1475 }
1476
blackman_window(int64_t window_length,bool periodic,std::optional<ScalarType> dtype_opt,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1477 Tensor blackman_window(
1478 int64_t window_length,
1479 bool periodic,
1480 std::optional<ScalarType> dtype_opt,
1481 std::optional<Layout> layout,
1482 std::optional<Device> device,
1483 std::optional<bool> pin_memory) {
1484 // See [Note: hacky wrapper removal for TensorOptions]
1485 ScalarType dtype = c10::dtype_or_default(dtype_opt);
1486 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
1487
1488 window_function_checks("blackman_window", options, window_length);
1489 if (window_length == 0) {
1490 return at::empty({0}, options);
1491 }
1492 if (window_length == 1) {
1493 return native::ones({1}, dtype, layout, device, pin_memory);
1494 }
1495 if (periodic) {
1496 window_length += 1;
1497 }
1498 // from https://en.wikipedia.org/wiki/Window_function#Blackman_window
1499 auto window =
1500 native::arange(window_length, dtype, layout, device, pin_memory)
1501 .mul_(c10::pi<double> / static_cast<double>(window_length - 1));
1502 window = window.mul(4).cos_().mul_(0.08) - window.mul(2).cos_().mul_(0.5) + 0.42;
1503 return periodic ? window.narrow(0, 0, window_length - 1) : std::move(window);
1504 }
1505
1506 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ hamming_window ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1507
hamming_window(int64_t window_length,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1508 Tensor hamming_window(int64_t window_length,
1509 std::optional<ScalarType> dtype,
1510 std::optional<Layout> layout,
1511 std::optional<Device> device,
1512 std::optional<bool> pin_memory) {
1513 return native::hamming_window(
1514 window_length, /*periodic=*/true, dtype, layout, device, pin_memory);
1515 }
1516
hamming_window(int64_t window_length,bool periodic,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1517 Tensor hamming_window(
1518 int64_t window_length,
1519 bool periodic,
1520 std::optional<ScalarType> dtype,
1521 std::optional<Layout> layout,
1522 std::optional<Device> device,
1523 std::optional<bool> pin_memory) {
1524 return native::hamming_window(
1525 window_length,
1526 periodic,
1527 /*alpha=*/0.54,
1528 dtype,
1529 layout,
1530 device,
1531 pin_memory);
1532 }
1533
hamming_window(int64_t window_length,bool periodic,double alpha,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1534 Tensor hamming_window(
1535 int64_t window_length,
1536 bool periodic,
1537 double alpha,
1538 std::optional<ScalarType> dtype,
1539 std::optional<Layout> layout,
1540 std::optional<Device> device,
1541 std::optional<bool> pin_memory) {
1542 return native::hamming_window(
1543 window_length, periodic, alpha, /*beta=*/0.46, dtype, layout, device, pin_memory);
1544 }
1545
hamming_window(int64_t window_length,bool periodic,double alpha,double beta,std::optional<ScalarType> dtype_opt,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1546 Tensor hamming_window(
1547 int64_t window_length,
1548 bool periodic,
1549 double alpha,
1550 double beta,
1551 std::optional<ScalarType> dtype_opt,
1552 std::optional<Layout> layout,
1553 std::optional<Device> device,
1554 std::optional<bool> pin_memory) {
1555 // See [Note: hacky wrapper removal for TensorOptions]
1556 ScalarType dtype = c10::dtype_or_default(dtype_opt);
1557 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
1558
1559 window_function_checks("hamming_window", options, window_length);
1560 if (window_length == 0) {
1561 return at::empty({0}, options);
1562 }
1563 if (window_length == 1) {
1564 return native::ones({1}, dtype, layout, device, pin_memory);
1565 }
1566 if (periodic) {
1567 window_length += 1;
1568 }
1569 auto window = native::arange(window_length, dtype, layout, device, pin_memory);
1570 window.mul_(c10::pi<double> * 2. / static_cast<double>(window_length - 1)).cos_().mul_(-beta).add_(alpha);
1571 return periodic ? window.narrow(0, 0, window_length - 1) : std::move(window);
1572 }
1573
1574 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ hann_window ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1575
hann_window(int64_t window_length,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1576 Tensor hann_window(int64_t window_length,
1577 std::optional<ScalarType> dtype,
1578 std::optional<Layout> layout,
1579 std::optional<Device> device,
1580 std::optional<bool> pin_memory) {
1581 return native::hann_window(window_length, /*periodic=*/true, dtype, layout, device, pin_memory);
1582 }
1583
hann_window(int64_t window_length,bool periodic,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1584 Tensor hann_window(
1585 int64_t window_length,
1586 bool periodic,
1587 std::optional<ScalarType> dtype,
1588 std::optional<Layout> layout,
1589 std::optional<Device> device,
1590 std::optional<bool> pin_memory) {
1591 // See [Note: hacky wrapper removal for TensorOptions]
1592 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
1593
1594 window_function_checks("hann_window", options, window_length);
1595 return native::hamming_window(
1596 window_length, periodic, /*alpha=*/0.5, /*beta=*/0.5, dtype, layout, device, pin_memory);
1597 }
1598
1599 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ kaiser_window ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1600
kaiser_window(int64_t window_length,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1601 Tensor kaiser_window(int64_t window_length,
1602 std::optional<ScalarType> dtype,
1603 std::optional<Layout> layout,
1604 std::optional<Device> device,
1605 std::optional<bool> pin_memory) {
1606 return native::kaiser_window(
1607 window_length,
1608 /*periodic=*/true,
1609 /*beta=*/12.0,
1610 dtype,
1611 layout,
1612 device,
1613 pin_memory);
1614 }
1615
kaiser_window(int64_t window_length,bool periodic,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1616 Tensor kaiser_window(int64_t window_length, bool periodic,
1617 std::optional<ScalarType> dtype,
1618 std::optional<Layout> layout,
1619 std::optional<Device> device,
1620 std::optional<bool> pin_memory) {
1621 return native::kaiser_window(window_length, periodic, /*beta=*/12.0, dtype, layout, device, pin_memory);
1622 }
1623
kaiser_window(int64_t window_length,bool periodic,double beta,std::optional<ScalarType> dtype_opt,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1624 Tensor kaiser_window(
1625 int64_t window_length,
1626 bool periodic,
1627 double beta,
1628 std::optional<ScalarType> dtype_opt,
1629 std::optional<Layout> layout,
1630 std::optional<Device> device,
1631 std::optional<bool> pin_memory) {
1632 // See [Note: hacky wrapper removal for TensorOptions]
1633 ScalarType dtype = c10::dtype_or_default(dtype_opt);
1634 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
1635
1636 window_function_checks("kaiser_window", options, window_length);
1637 // short-circuit for `meta`.
1638 if (device == kMeta) {
1639 return at::empty({window_length}, options);
1640 }
1641
1642 if (window_length == 0) {
1643 return at::empty({0}, options);
1644 }
1645 if (window_length == 1) {
1646 return at::ones({1}, options);
1647 }
1648 if (periodic) {
1649 window_length += 1;
1650 }
1651 auto initial = at::arange(window_length, options);
1652 auto window = at::empty(window_length, options);
1653 auto iter = TensorIterator::unary_op(window, initial);
1654 kaiser_window_stub(iter.device_type(), iter, window_length, beta);
1655 return periodic ? window.narrow(0, 0, window_length - 1) : std::move(window);
1656 }
1657
1658 // ~~~~~~~~~~~~~~~~~~~~~~~~~~ vandermonde_matrix ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1659
1660
vander(const Tensor & x,std::optional<int64_t> N,bool increasing)1661 Tensor vander(const Tensor& x, std::optional<int64_t> N, bool increasing) {
1662 TORCH_CHECK(x.dim() == 1, "x must be a one-dimensional tensor.");
1663
1664 // Acquires n, defaulting to size if not provided
1665 int64_t n = x.size(0);
1666 if (N.has_value()) {
1667 n = *N;
1668 TORCH_CHECK(n >= 0, "N must be non-negative.");
1669 }
1670
1671 // Note: result is long if x is an integer tensor (like int8) because
1672 // cumprod promotes integer tensors to long
1673 auto result = at::empty({x.size(0), n}, x.options().dtype(at::promote_types(x.scalar_type(), c10::ScalarType::Long)));
1674
1675 if (n > 0) {
1676 result.select(1, 0).fill_(1);
1677 }
1678 if (n > 1) {
1679 result.slice(1, 1).copy_(x.unsqueeze(1));
1680 result.slice(1, 1).copy_(at::cumprod(result.slice(1, 1), 1));
1681 }
1682
1683 if (!increasing) {
1684 return at::flip(result, {1});
1685 }
1686 return result;
1687 }
1688
1689 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ tensor ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1690
1691 template <typename T>
tensor_cpu(ArrayRef<T> values,const TensorOptions & options)1692 Tensor tensor_cpu(ArrayRef<T> values, const TensorOptions& options) {
1693 return at::detail::tensor_cpu(values, options);
1694 }
1695
1696 template <typename T>
tensor_backend(ArrayRef<T> values,const TensorOptions & options)1697 Tensor tensor_backend(ArrayRef<T> values, const TensorOptions& options) {
1698 return at::detail::tensor_backend(values, options);
1699 }
1700
1701 template <typename T>
tensor_complex_cpu(ArrayRef<T> values,const TensorOptions & options)1702 Tensor tensor_complex_cpu(ArrayRef<T> values, const TensorOptions& options) {
1703 return at::detail::tensor_complex_cpu(values, options);
1704 }
1705
1706 template <typename T>
tensor_complex_backend(ArrayRef<T> values,const TensorOptions & options)1707 Tensor tensor_complex_backend(ArrayRef<T> values, const TensorOptions& options) {
1708 return at::detail::tensor_complex_backend(values, options);
1709 }
1710
from_file(c10::string_view filename,std::optional<bool> shared,std::optional<int64_t> size,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1711 Tensor from_file(c10::string_view filename, std::optional<bool> shared, std::optional<int64_t> size,
1712 std::optional<ScalarType> dtype,
1713 std::optional<Layout> layout,
1714 std::optional<Device> device,
1715 std::optional<bool> pin_memory) {
1716 // See [Note: hacky wrapper removal for TensorOptions]
1717 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
1718
1719 TORCH_CHECK(!options.pinned_memory(), "tensors constructed from a file cannot be pinned");
1720 int64_t my_size = size.value_or(0);
1721 int flags = shared.value_or(false) ? ALLOCATOR_MAPPED_SHARED : 0;
1722 auto my_dtype = options.dtype();
1723 size_t size_bytes = my_size * my_dtype.itemsize();
1724 auto storage_impl = c10::make_intrusive<at::StorageImpl>(
1725 c10::StorageImpl::use_byte_size_t(),
1726 size_bytes,
1727 MapAllocator::makeDataPtr(
1728 std::string(filename), flags, size_bytes, nullptr),
1729 /*allocator=*/nullptr,
1730 /*resizable=*/false);
1731 auto tensor = detail::make_tensor<at::TensorImpl>(
1732 storage_impl, at::DispatchKey::CPU, my_dtype);
1733 tensor.unsafeGetTensorImpl()->set_sizes_contiguous({my_size});
1734 return tensor;
1735 }
1736
1737 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ clone ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1738
clone(const Tensor & src,std::optional<c10::MemoryFormat> optional_memory_format)1739 Tensor clone(const Tensor& src, std::optional<c10::MemoryFormat> optional_memory_format) {
1740 auto memory_format =
1741 optional_memory_format.value_or(MemoryFormat::Preserve);
1742 Tensor self;
1743 if (memory_format == MemoryFormat::Preserve) {
1744 if (src.is_non_overlapping_and_dense()) {
1745 // Copy all strides, this is marginally faster than calling empty_like
1746 self = at::empty_strided_symint(src.sym_sizes(), src.sym_strides(), src.options());
1747 } else {
1748 self = at::empty_like(src);
1749 }
1750 } else {
1751 self = at::empty_like(src, src.options(), memory_format);
1752 }
1753
1754 if (src._is_zerotensor()) {
1755 self.zero_();
1756 } else {
1757 self.copy_(src);
1758 }
1759 return self;
1760 }
1761
1762 // ~~~~~~~~~~~~~~~~~~~~~~~~~ named tensor overloads ~~~~~~~~~~~~~~~~~~~~~~~~~~~
1763 // In the short term, these exist.
1764 // In the long term, we should move DimnameList into TensorOptions to avoid
1765 // having these overloads.
1766
full(IntArrayRef size,const Scalar & fill_value,std::optional<DimnameList> names,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1767 Tensor full(
1768 IntArrayRef size,
1769 const Scalar& fill_value,
1770 std::optional<DimnameList> names,
1771 std::optional<ScalarType> dtype,
1772 std::optional<Layout> layout,
1773 std::optional<Device> device,
1774 std::optional<bool> pin_memory) {
1775 // See [Note: hacky wrapper removal for TensorOptions]
1776 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
1777
1778
1779 TORCH_CHECK(options.layout() != kSparse,
1780 "full(...) is not implemented for sparse layout");
1781
1782 auto result = at::empty(size, names, infer_full_options(fill_value, options));
1783 return result.fill_(fill_value);
1784 }
1785
ones(IntArrayRef size,std::optional<DimnameList> names,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1786 Tensor ones(
1787 IntArrayRef size,
1788 std::optional<DimnameList> names,
1789 std::optional<ScalarType> dtype,
1790 std::optional<Layout> layout,
1791 std::optional<Device> device,
1792 std::optional<bool> pin_memory) {
1793 // See [Note: hacky wrapper removal for TensorOptions]
1794
1795 return native::full(
1796 size, /*fill_value=*/1., names, dtype, layout, device, pin_memory);
1797 }
1798
zeros(IntArrayRef size,std::optional<DimnameList> names,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1799 Tensor zeros(
1800 IntArrayRef size,
1801 std::optional<DimnameList> names,
1802 std::optional<ScalarType> dtype,
1803 std::optional<Layout> layout,
1804 std::optional<Device> device,
1805 std::optional<bool> pin_memory) {
1806 return native::full(size, /*fill_value=*/0., names, dtype, layout, device, pin_memory);
1807 }
1808
randn(IntArrayRef size,std::optional<DimnameList> names,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1809 Tensor randn(
1810 IntArrayRef size,
1811 std::optional<DimnameList> names,
1812 std::optional<ScalarType> dtype,
1813 std::optional<Layout> layout,
1814 std::optional<Device> device,
1815 std::optional<bool> pin_memory) {
1816 return native::randn(size, std::nullopt, names, dtype, layout, device, pin_memory);
1817 }
1818
randn(IntArrayRef size,std::optional<Generator> generator,std::optional<DimnameList> names,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1819 Tensor randn(
1820 IntArrayRef size,
1821 std::optional<Generator> generator,
1822 std::optional<DimnameList> names,
1823 std::optional<ScalarType> dtype,
1824 std::optional<Layout> layout,
1825 std::optional<Device> device,
1826 std::optional<bool> pin_memory) {
1827 // See [Note: hacky wrapper removal for TensorOptions]
1828 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
1829
1830 auto result = at::empty(size, names, options);
1831 return result.normal_(0, 1, std::move(generator));
1832 }
1833
rand(IntArrayRef size,std::optional<DimnameList> names,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1834 Tensor rand(
1835 IntArrayRef size,
1836 std::optional<DimnameList> names,
1837 std::optional<ScalarType> dtype,
1838 std::optional<Layout> layout,
1839 std::optional<Device> device,
1840 std::optional<bool> pin_memory) {
1841 return native::rand(size, std::nullopt, names, dtype, layout, device, pin_memory);
1842 }
1843
rand(IntArrayRef size,std::optional<Generator> generator,std::optional<DimnameList> names,std::optional<ScalarType> dtype,std::optional<Layout> layout,std::optional<Device> device,std::optional<bool> pin_memory)1844 Tensor rand(
1845 IntArrayRef size,
1846 std::optional<Generator> generator,
1847 std::optional<DimnameList> names,
1848 std::optional<ScalarType> dtype,
1849 std::optional<Layout> layout,
1850 std::optional<Device> device,
1851 std::optional<bool> pin_memory) {
1852 // See [Note: hacky wrapper removal for TensorOptions]
1853 TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
1854
1855 auto result = at::empty(size, names, options);
1856 return result.uniform_(0, 1, std::move(generator));
1857 }
1858
1859
1860 DEFINE_DISPATCH(kaiser_window_stub);
1861
1862 } // namespace at::native
1863