xref: /aosp_15_r20/external/pytorch/aten/src/ATen/native/cpu/WeightNormKernel.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1 #pragma once
2 #include <ATen/native/DispatchStub.h>
3 #include <cstdint>
4 
5 namespace at {
6 class TensorBase;
7 }
8 
9 namespace at::native {
10 
11 using weight_norm_fn = void(*)(
12     TensorBase&, TensorBase&, const TensorBase&, const TensorBase&, int64_t);
13 using weight_norm_backward_fn = void(*)(
14     TensorBase&, TensorBase&, const TensorBase&, const TensorBase&,
15     const TensorBase&, const TensorBase&, int64_t);
16 
17 DECLARE_DISPATCH(weight_norm_fn, weight_norm_stub);
18 DECLARE_DISPATCH(weight_norm_backward_fn, weight_norm_backward_stub);
19 
20 } // namespace at::native
21