mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
Add matches_jit_signature to recent native functions
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/17805 Differential Revision: D14388004 Pulled By: cpuhrsch fbshipit-source-id: c50580b6fe1e9cfefed91aaa526376325d9f9c0d
This commit is contained in:
committed by
Facebook Github Bot
parent
fe90ee9dc8
commit
11f50e73e3
@ -1608,14 +1608,17 @@
|
||||
CUDA: batch_norm_cuda
|
||||
|
||||
- func: batch_norm_stats(Tensor input, float eps) -> (Tensor, Tensor)
|
||||
matches_jit_signature: True
|
||||
dispatch:
|
||||
CUDA: batch_norm_stats_cuda
|
||||
|
||||
- func: batch_norm_elemt(Tensor input, Tensor? weight, Tensor? bias, Tensor mean, Tensor invstd, float eps) -> Tensor
|
||||
matches_jit_signature: True
|
||||
dispatch:
|
||||
CUDA: batch_norm_elemt_cuda
|
||||
|
||||
- func: batch_norm_gather_stats(Tensor input, Tensor mean, Tensor invstd, Tensor? running_mean, Tensor? running_var, float momentum, float eps, int count) -> (Tensor, Tensor)
|
||||
matches_jit_signature: True
|
||||
dispatch:
|
||||
CUDA: batch_norm_gather_stats_cuda
|
||||
|
||||
@ -1626,10 +1629,12 @@
|
||||
CUDA: batch_norm_backward_cuda
|
||||
|
||||
- func: batch_norm_backward_reduce(Tensor grad_out, Tensor input, Tensor mean, Tensor invstd, bool input_g, bool weight_g, bool bias_g) -> (Tensor, Tensor, Tensor, Tensor)
|
||||
matches_jit_signature: True
|
||||
dispatch:
|
||||
CUDA: batch_norm_backward_reduce_cuda
|
||||
|
||||
- func: batch_norm_backward_elemt(Tensor grad_out, Tensor input, Tensor mean, Tensor invstd, Tensor? weight, Tensor mean_dy, Tensor mean_dy_xmu) -> Tensor
|
||||
matches_jit_signature: True
|
||||
dispatch:
|
||||
CUDA: batch_norm_backward_elemt_cuda
|
||||
|
||||
|
Reference in New Issue
Block a user