Fix Native signature for optional Tensor arguments (#50767)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/50767
The native signature for optional tensor arguments wrongly produced "Tensor" instead of "optional<Tensor>". We didn't notice this because all internal ops currently use hacky_wrapper, and for hacky_wrapper, "Tensor" is correct.
This PR fixes that and ports one op (batch_norm) to not use hacky_wrapper anymore as a proof of fix.
ghstack-source-id: 120017543
Test Plan: waitforsandcastle
Reviewed By: bhosmer
Differential Revision: D25960941
fbshipit-source-id: ca6fe133109b5d85cff52390792cf552f12d9590
diff --git a/aten/src/ATen/native/Normalization.cpp b/aten/src/ATen/native/Normalization.cpp
index ea4a54c..91dafc1 100644
--- a/aten/src/ATen/native/Normalization.cpp
+++ b/aten/src/ATen/native/Normalization.cpp
@@ -490,9 +490,13 @@
}
Tensor batch_norm(
- const Tensor& input, const Tensor& weight /* optional */, const Tensor& bias /* optional */,
- const Tensor& running_mean /* optional */, const Tensor& running_var /* optional */,
+ const Tensor& input, const c10::optional<Tensor>& weight_opt, const c10::optional<Tensor>& bias_opt,
+ const c10::optional<Tensor>& running_mean_opt, const c10::optional<Tensor>& running_var_opt,
bool training, double momentum, double eps, bool cudnn_enabled) {
+ const Tensor& weight = c10::value_or_else(weight_opt, [] {return Tensor();});
+ const Tensor& bias = c10::value_or_else(bias_opt, [] {return Tensor();});
+ const Tensor& running_mean = c10::value_or_else(running_mean_opt, [] {return Tensor();});
+ const Tensor& running_var = c10::value_or_else(running_var_opt, [] {return Tensor();});
if (input.numel()==0){
//don't return view of input, don't return empty tensor because it will break gradient chain
auto out = input.clone();
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
index 1c329af..91b4dd7 100644
--- a/aten/src/ATen/native/native_functions.yaml
+++ b/aten/src/ATen/native/native_functions.yaml
@@ -697,7 +697,6 @@
use_c10_dispatcher: hacky_wrapper_for_legacy_signatures
- func: batch_norm(Tensor input, Tensor? weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool training, float momentum, float eps, bool cudnn_enabled) -> Tensor
- use_c10_dispatcher: hacky_wrapper_for_legacy_signatures
- func: quantized_batch_norm(Tensor input, Tensor? weight, Tensor? bias, Tensor mean, Tensor var, float eps, float output_scale, int output_zero_point) -> Tensor
use_c10_dispatcher: hacky_wrapper_for_legacy_signatures
diff --git a/tools/codegen/api/native.py b/tools/codegen/api/native.py
index af82210..9afae68 100644
--- a/tools/codegen/api/native.py
+++ b/tools/codegen/api/native.py
@@ -29,10 +29,13 @@
def argumenttype_type(t: Type, *, mutable: bool, binds: ArgName) -> CType:
if str(t) == 'Tensor?':
+ tensor_type: CType = BaseCType('Tensor', binds)
+ if local.use_c10_dispatcher() is not UseC10Dispatcher.hacky_wrapper_for_legacy_signatures:
+ tensor_type = OptionalCType(tensor_type)
if mutable:
- return MutRefCType(BaseCType('Tensor', binds))
+ return MutRefCType(tensor_type)
else:
- return ConstRefCType(BaseCType('Tensor', binds))
+ return ConstRefCType(tensor_type)
elif str(t) == 'Tensor?[]':
return BaseCType('const c10::List<c10::optional<Tensor>> &', binds)
return cpp.argumenttype_type(t, mutable=mutable, binds=binds)
