Propagate NaN through threshold (#10277)

ssnl · facebook-github-bot · commit 254dedf6043d · 2018-08-15T12:59:31.000-07:00
Summary: Fixes #10238 Pull Request resolved: #10277 Reviewed By: SsnL Differential Revision: D9199825 Pulled By: soumith fbshipit-source-id: 8ee7f9a72d9546d429f311c3f6028461d3c93fe2
diff --git a/aten/src/THCUNN/Threshold.cu b/aten/src/THCUNN/Threshold.cu
@@ -17,7 +17,7 @@ struct ThresholdUpdateOutput
   __device__ __forceinline__ void operator()(T *out, T *in)
   {
     T x = *in;
-    *out = (x > threshold_) ? x : val_;
+    *out = (x <= threshold_) ? val_ : x;  // this order propagates NaN
   }
 };
 
@@ -35,7 +35,7 @@ struct ThresholdUpdateOutputIP
 
   __device__ __forceinline__ void operator()(T *x)
   {
-    *x = (*x > threshold_) ? *x : val_;
+    *x = (*x <= threshold_) ? val_ : *x;  // this order propagates NaN
   }
 };
 
@@ -51,7 +51,7 @@ struct ThresholdUpdateGradInput
   __device__ __forceinline__ void operator()(
     T *gradInput, T *input, T *gradOutput) const
   {
-    *gradInput = (*input > threshold_) ? *gradOutput : ScalarConvert<int, T>::to(0);
+    *gradInput = (*input <= threshold_) ? ScalarConvert<int, T>::to(0) : *gradOutput;  // this order propagates NaN
   }
 };
 
@@ -67,7 +67,7 @@ struct ThresholdUpdateGradInputIP
   __device__ __forceinline__ void operator()(
     T *gradOutput, T *input) const
   {
-    *gradOutput = (*input > threshold_) ? *gradOutput : ScalarConvert<int, T>::to(0);
+    *gradOutput = (*input <= threshold_) ? ScalarConvert<int, T>::to(0) : *gradOutput;  // this order propagates NaN
   }
 };
 
diff --git a/aten/src/THNN/generic/Threshold.c b/aten/src/THNN/generic/Threshold.c
@@ -24,7 +24,7 @@ void THNN_(Threshold_updateOutput)(
   {
     THTensor_(resizeAs)(output, input);
     TH_TENSOR_APPLY2(real, output, real, input,
-      *output_data = (*input_data > threshold) ? *input_data : val;
+      *output_data = (*input_data <= threshold) ? val : *input_data;  // this order propagates NaN
     );
   }
 }
@@ -52,10 +52,10 @@ void THNN_(Threshold_updateGradInput)(
   {
     THTensor_(resizeAs)(gradInput, input);
     TH_TENSOR_APPLY3(real, gradInput, real, gradOutput, real, input,
-      if ((*input_data) > threshold)
-        *gradInput_data = *gradOutput_data;
-      else
+      if ((*input_data) <= threshold)
         *gradInput_data = 0;
+      else
+        *gradInput_data = *gradOutput_data;  // let NaN case pass through here
     );
   }
 }
diff --git a/test/test_nn.py b/test/test_nn.py
@@ -1614,6 +1614,8 @@ def test(nonlinearity, *args, **kwargs):
         test('softmax', 0)
         test('log_softmax', 0)
         test('leaky_relu', 0.2)
+        test('threshold', 3, 2)
+        test('threshold', 3, 2, inplace=True)
 
     def test_nonlinearity_propagate_nan(self):
         self._test_nonlinearity_propagate_nan('cpu')

Original file line number	Diff line number	Diff line change
`@@ -17,7 +17,7 @@ struct ThresholdUpdateOutput`
`17`	`17`	`__device__ __forceinline__ void operator()(T out, T in)`
`18`	`18`	`{`
`19`	`19`	`T x = *in;`
`20`		`- *out = (x > threshold_) ? x : val_;`
	`20`	`+ *out = (x <= threshold_) ? val_ : x; // this order propagates NaN`
`21`	`21`	`}`
`22`	`22`	`};`
`23`	`23`
`@@ -35,7 +35,7 @@ struct ThresholdUpdateOutputIP`
`35`	`35`
`36`	`36`	`__device__ __forceinline__ void operator()(T *x)`
`37`	`37`	`{`
`38`		`- x = (x > threshold_) ? *x : val_;`
	`38`	`+ x = (x <= threshold_) ? val_ : *x; // this order propagates NaN`
`39`	`39`	`}`
`40`	`40`	`};`
`41`	`41`
`@@ -51,7 +51,7 @@ struct ThresholdUpdateGradInput`
`51`	`51`	`__device__ __forceinline__ void operator()(`
`52`	`52`	`T gradInput, T input, T *gradOutput) const`
`53`	`53`	`{`
`54`		`- gradInput = (input > threshold_) ? *gradOutput : ScalarConvert<int, T>::to(0);`
	`54`	`+ gradInput = (input <= threshold_) ? ScalarConvert<int, T>::to(0) : *gradOutput; // this order propagates NaN`
`55`	`55`	`}`
`56`	`56`	`};`
`57`	`57`
`@@ -67,7 +67,7 @@ struct ThresholdUpdateGradInputIP`
`67`	`67`	`__device__ __forceinline__ void operator()(`
`68`	`68`	`T gradOutput, T input) const`
`69`	`69`	`{`
`70`		`- gradOutput = (input > threshold_) ? *gradOutput : ScalarConvert<int, T>::to(0);`
	`70`	`+ gradOutput = (input <= threshold_) ? ScalarConvert<int, T>::to(0) : *gradOutput; // this order propagates NaN`
`71`	`71`	`}`
`72`	`72`	`};`
`73`	`73`
Original file line number	Diff line number	Diff line change
`@@ -24,7 +24,7 @@ void THNN_(Threshold_updateOutput)(`
`24`	`24`	`{`
`25`	`25`	`THTensor_(resizeAs)(output, input);`
`26`	`26`	`TH_TENSOR_APPLY2(real, output, real, input,`
`27`		`- output_data = (input_data > threshold) ? *input_data : val;`
	`27`	`+ output_data = (input_data <= threshold) ? val : *input_data; // this order propagates NaN`
`28`	`28`	`);`
`29`	`29`	`}`
`30`	`30`	`}`
`@@ -52,10 +52,10 @@ void THNN_(Threshold_updateGradInput)(`
`52`	`52`	`{`
`53`	`53`	`THTensor_(resizeAs)(gradInput, input);`
`54`	`54`	`TH_TENSOR_APPLY3(real, gradInput, real, gradOutput, real, input,`
`55`		`- if ((*input_data) > threshold)`
`56`		`- gradInput_data = gradOutput_data;`
`57`		`- else`
	`55`	`+ if ((*input_data) <= threshold)`
`58`	`56`	`*gradInput_data = 0;`
	`57`	`+ else`
	`58`	`+ gradInput_data = gradOutput_data; // let NaN case pass through here`
`59`	`59`	`);`
`60`	`60`	`}`
`61`	`61`	`}`