xref: /aosp_15_r20/external/pytorch/aten/src/ATen/native/cpu/StackKernel.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 // Copyright 2004-present Facebook. All Rights Reserved.
2 #define TORCH_ASSERT_ONLY_METHOD_OPERATORS
3 #include <ATen/core/Tensor.h>
4 
5 #include <ATen/Dispatch.h>
6 #include <ATen/native/cpu/StackKernel.h>
7 #include <ATen/native/cpu/SerialStackImpl.h>
8 
9 namespace at::native {
10 
11 namespace {
12 
stack_serial_kernel(Tensor & result,TensorList tensors,int64_t dim)13 void stack_serial_kernel(Tensor& result, TensorList tensors, int64_t dim) {
14   AT_DISPATCH_FLOATING_TYPES(
15       result.scalar_type(), "stack_serial_kernel", [&]() {
16         detail::stack_serial_kernel_impl<scalar_t, TensorList>(result, tensors, dim);
17       });
18 }
19 
20 } // anonymous namespace
21 
22 REGISTER_DISPATCH(stack_serial_stub, &stack_serial_kernel);
23 
24 } // namespace at::native
25