Fix the clamp special case and gradient problem on None, add None to JIT (pytorch#9596)

Summary:
Supersedes pytorch#8925

This PR fixes pytorch#8502, it fixes the gradients problem for clamp when passing None to the function, and add support for the NoneLiteral and NoneType in script to enable clamp tests. Now we could have corner cases like:

```python
torch.jit.script
def func():
    x = torch.randn(3, 3, requires_grad=True)
    y = torch.clamp(x, None, 0) # max = 0
    y = torch.clamp(x, min=None, max=0)
```

In both JIT and Aten, we use Scalar(NAN) as a sentinel value when passing None type to function clamp, this is the current way we used to support None type in JIT and to solve the gradient problem when user explicitly passing None into clamp.

In JIT side, we create a tensor(NAN) and undefinedTensor if we encounter None when matching the function schema, and later in the interpreter, it will translate to Scalar(NAN) if needed.

Ideally we don't need clamp_min and clamp_max in ATenNative/Autograd and could only support clamp after this change, but since bunch of other operators (e.g. Activation.cpp, Loss.cpp) is using clamp_min in several places, we will still have the functions available, but all python invocations will only call clamp instead of clamp_min/max (with calling underlying th_max/th_min in clamp).

zdevito jamesr66a
Pull Request resolved: pytorch#9596

Reviewed By: zdevito

Differential Revision: D8940839

Pulled By: wanchaol

fbshipit-source-id: c543a867b82e0ab8c99384773b173fdde2605d28

Author: wanchaol

aten/src/ATen/native/UnaryOps.cpp

@@ -46,7 +46,15 @@ Tensor clamp_min(const Tensor& self, Scalar min) {
 }
 
 Tensor& _clamp__cpu(Tensor& self, Scalar min, Scalar max) {
-  return _th_clamp_(self, min, max);
+  if (!std::isnan(min.toDouble()) && !std::isnan(max.toDouble())) {
+    return _th_clamp_(self, min, max);
+  } else if (std::isnan(min.toDouble())) {
+    return _th_clamp_max_(self, max);
+  } else if (std::isnan(max.toDouble())) {
+    return _th_clamp_min_(self, min);
+  } else {
+    return self;
+  }
 }
 
 Tensor& _clamp_out_cpu(
@@ -56,7 +64,14 @@ Tensor& _clamp_out_cpu(
     Scalar max) {
   result.resize_(self.sizes());
   result.copy_(self);
-  return _th_clamp_(result, min, max);
+  if (!std::isnan(min.toDouble()) && !std::isnan(max.toDouble())) {
+    _th_clamp_(result, min, max);
+  } else if (std::isnan(min.toDouble())) {
+    _th_clamp_max_(result, max);
+  } else if (std::isnan(max.toDouble())) {
+    _th_clamp_min_(result, min);
+  }
+  return result;
 }
 
 Tensor& _clamp_max__cpu(Tensor& self, Scalar max) {

aten/src/ATen/native/cuda/CUDAUnaryOps.cpp

@@ -3,7 +3,15 @@
 namespace at { namespace native {
 
 Tensor& _clamp__cuda(Tensor& self, Scalar min, Scalar max) {
-  return _th_clamp_(self, min, max);
+  if (!std::isnan(min.toDouble()) && !std::isnan(max.toDouble())) {
+    return _th_clamp_(self, min, max);
+  } else if (std::isnan(min.toDouble())) {
+    return _th_clamp_max_(self, max);
+  } else if (std::isnan(max.toDouble())) {
+    return _th_clamp_min_(self, min);
+  } else {
+    return self;
+  }
 }
 
 Tensor& _clamp_out_cuda(
@@ -13,7 +21,14 @@ Tensor& _clamp_out_cuda(
     Scalar max) {
   result.resize_(self.sizes());
   result.copy_(self);
-  return _th_clamp_(result, min, max);
+  if (!std::isnan(min.toDouble()) && !std::isnan(max.toDouble())) {
+    _th_clamp_(result, min, max);
+  } else if (std::isnan(min.toDouble())) {
+    _th_clamp_max_(result, max);
+  } else if (std::isnan(max.toDouble())) {
+    _th_clamp_min_(result, min);
+  }
+  return result;
 }
 
 Tensor& _clamp_max__cuda(Tensor& self, Scalar max) {

aten/src/ATen/native/native_functions.yaml

@@ -266,17 +266,26 @@
 - func: chunk(Tensor self, int64_t chunks, int64_t dim=0) -> TensorList
 
 - func: clamp(Tensor self, Scalar min, Scalar max) -> Tensor
+  python_default_init:
+    min: NAN
+    max: NAN
 
 - func: clamp_(Tensor self, Scalar min, Scalar max) -> Tensor
   dispatch:
     CPU: _clamp__cpu
     CUDA: _clamp__cuda
+  python_default_init:
+    min: NAN
+    max: NAN
 
 - func: clamp_out(Tensor result, Tensor self, Scalar min, Scalar max) -> Tensor
   variants: function
   dispatch:
     CPU: _clamp_out_cpu
     CUDA: _clamp_out_cuda
+  python_default_init:
+    min: NAN
+    max: NAN
 
 - func: clamp_max(Tensor self, Scalar max) -> Tensor
 

test/expect/TestScript.test_python_frontend.expect

@@ -5,6 +5,10 @@
     (param (ident y) (tensor_type))
     (param (ident z) (tensor_type)))
   (list
+    (assign
+      (list (variable (ident q)))
+      (=)
+      (None))
     (assign
       (list (variable (ident q)))
       (=)

test/test_jit.py

@@ -1629,6 +1629,7 @@ def to_int(x):
 
     def test_python_frontend(self):
         def fn(x, y, z):
+            q = None
             q = x + y - z.sigmoid()
             print(q)
             w = -z
@@ -1862,6 +1863,28 @@ def test_script_for_in_range_if_ast(x):
 
         self.assertEqual(test_script_for_in_range_if_ast(*inputs).shape[0], 20)
 
+    def test_script_None(self):
+        def func(x):
+            output = None
+            output = x
+            return output
+
+        self.checkScript(func, [torch.arange(0, 2)], optimize=True)
+
+    def test_script_clamp_none(self):
+        # TODO: could not enable default/optional argument for None in JIT
+        # result from Aten native python_default_init for clamp, it is used
+        # in Aten but not in JIT, need to fix type/default arg system in ATen
+        def test_script_clamp_max_none(x):
+            return torch.clamp(x, min=None, max=2)
+
+        def test_script_clamp_min_none(x):
+            return torch.clamp(x, min=2, max=None)
+
+        input = [torch.arange(0, 3)]
+        self.checkScript(test_script_clamp_max_none, input, optimize=True)
+        self.checkScript(test_script_clamp_min_none, input, optimize=True)
+
     def test_script_bool_constant(self):
         script = '''
         def test_script_bool_constant():
@@ -4845,10 +4868,6 @@ def forward(self, x, y):
 
 # known to be failing in script
 EXCLUDE_SCRIPT = {
-    'test_clamp_max',
-    'test_clamp_max_scalar',
-    'test_clamp_min',
-    'test_clamp_min_scalar',
     # TODO: Fix var/std
     # there are two schemas for var (and std):
     # (1) var(Tensor, int, *, bool, bool, Tensor)

tools/autograd/derivatives.yaml

@@ -168,11 +168,10 @@
 - name: ceil(Tensor self)
   self: zeros_like(grad)
 
-# For clamp, clamp_min, and clamp_max, gradient is not defined at the
-# boundaries. But empirically it's helpful to be able to get gradient on min and
-# max, so we return the subgradient 1 for these cases.
+# For clamp, gradient is not defined at the boundaries. But empirically it's helpful 
+# to be able to get gradient on min and max, so we return the subgradient 1 for these cases.
 - name: clamp(Tensor self, Scalar min, Scalar max)
-  self: grad * ((self >= min) * (self <= max)).type_as(grad)
+  self: clamp_backward(grad, self, min, max)
 
 - name: clamp_min(Tensor self, Scalar min)
   self: grad * (self >= min).type_as(grad)

tools/autograd/gen_python_functions.py

@@ -17,7 +17,7 @@
 
 # These functions require manual Python bindings or are not exposed to Python
 SKIP_PYTHON_BINDINGS = [
-    'alias', 'contiguous', 'clamp.*', 'is_cuda', 'is_sparse', 'size', 'stride',
+    'alias', 'contiguous', 'is_cuda', 'is_sparse', 'size', 'stride',
     '.*_backward', '.*_backward_(out|input|weight|bias)', '.*_forward',
     '.*_forward_out', 'sparse_raw_resize_', '_unsafe_view', 'tensor',
     'sparse_coo_tensor', 'th_sparse_coo_tensor', 'native_sparse_coo_tensor',

tools/autograd/templates/Functions.cpp

@@ -431,6 +431,17 @@ std::vector<Tensor> cat_tensors_backward(const Tensor & grad, const std::vector<
   return grad_inputs;
 }
 
+Tensor clamp_backward(const Tensor & grad, const Tensor &self, const Scalar & min, const Scalar & max) {
+  // clamp: gradients not defined on min and max, so we return the subgradient 1 for these cases.
+  if (std::isnan(min.toFloat())) {
+    return grad * (self <= max).type_as(grad);
+  } else if (std::isnan(max.toFloat())) {
+    return grad * (self >= min).type_as(grad);
+  } else {
+    return grad * ((self >= min) * (self <= max)).type_as(grad);
+  }
+}
+
 Tensor mm_mat1_backward(const Tensor & grad, const Tensor & mat2, IntList sizes, IntList strides, const Scalar & alpha) {
   // if input was column-major, return grad as column-order for efficiency
   if (strides[0] == 1 && strides[1] == sizes[0]) {

tools/autograd/templates/python_torch_functions.cpp

@@ -191,41 +191,6 @@ static PyObject * THPVariable_as_tensor(PyObject* self, PyObject* args, PyObject
   END_HANDLE_TH_ERRORS
 }
 
-// The Python clamp() syntax has to be mapped to one of three C++ functions
-static PyObject * THPVariable_clamp(PyObject* module, PyObject* args, PyObject* kwargs)
-{
-  HANDLE_TH_ERRORS
-  static PythonArgParser parser({
-    "clamp(Tensor input, Scalar min=None, Scalar max=None, *, Tensor out=None)",
-  });
-
-  ParsedArgs<4> parsed_args;
-  auto r = parser.parse(args, kwargs, parsed_args);
-  if (!r.isNone(1) && !r.isNone(2)) {
-    if (!r.isNone(3)) {
-        return wrap(dispatch_clamp(r.tensor(0), r.scalar(1), r.scalar(2), r.tensor(3)));
-    } else {
-        return wrap(dispatch_clamp(r.tensor(0), r.scalar(1), r.scalar(2)));
-    }
-  } else if (!r.isNone(1)) {
-    if (!r.isNone(3)) {
-        return wrap(dispatch_clamp_min(r.tensor(0), r.scalar(1), r.tensor(3)));
-    } else {
-        return wrap(dispatch_clamp_min(r.tensor(0), r.scalar(1)));
-    }
-  } else if (!r.isNone(2)) {
-    if (!r.isNone(3)) {
-        return wrap(dispatch_clamp_max(r.tensor(0), r.scalar(2), r.tensor(3)));
-    } else {
-        return wrap(dispatch_clamp_max(r.tensor(0), r.scalar(2)));
-    }
-  } else {
-    throw std::runtime_error("At least one of 'min' or 'max' must not be None");
-  }
-  Py_RETURN_NONE;
-  END_HANDLE_TH_ERRORS
-}
-
 static PyObject * THPVariable_from_numpy(PyObject* module, PyObject* arg)
 {
   HANDLE_TH_ERRORS
@@ -271,7 +236,6 @@ static PyObject * THPVariable_tensor(PyObject* self, PyObject* args, PyObject* k
 static PyMethodDef torch_functions[] = {
   {"arange", (PyCFunction)THPVariable_arange, METH_VARARGS | METH_KEYWORDS | METH_STATIC, NULL},
   {"as_tensor", (PyCFunction)THPVariable_as_tensor, METH_VARARGS | METH_KEYWORDS | METH_STATIC, NULL},
-  {"clamp", (PyCFunction)THPVariable_clamp, METH_VARARGS | METH_KEYWORDS | METH_STATIC, NULL},
   {"dsmm", (PyCFunction)THPVariable_mm, METH_VARARGS | METH_KEYWORDS | METH_STATIC, NULL},
   {"from_numpy", (PyCFunction)THPVariable_from_numpy, METH_STATIC | METH_O, NULL},
   {"hsmm", (PyCFunction)THPVariable_hspmm, METH_VARARGS | METH_KEYWORDS | METH_STATIC, NULL},

tools/autograd/templates/python_torch_functions_dispatch.h

@@ -34,32 +34,6 @@ static void maybe_initialize_cuda(const at::Type &type) {
   }
 }
 
-// manual dispatch code for clamp
-inline Tensor dispatch_clamp(const Tensor & self, Scalar min, Scalar max) {
-  AutoNoGIL no_gil;
-  return self.clamp(min, max);
-}
-inline Tensor dispatch_clamp_min(const Tensor & self, Scalar min) {
-  AutoNoGIL no_gil;
-  return self.clamp_min(min);
-}
-inline Tensor dispatch_clamp_max(const Tensor & self, Scalar max) {
-  AutoNoGIL no_gil;
-  return self.clamp_max(max);
-}
-inline Tensor & dispatch_clamp(const Tensor & self, Scalar min, Scalar max, Tensor result) {
-  AutoNoGIL no_gil;
-  return at::clamp_out(result, self, min, max);
-}
-inline Tensor & dispatch_clamp_min(const Tensor & self, Scalar min, Tensor result) {
-  AutoNoGIL no_gil;
-  return at::clamp_min_out(result, self, min);
-}
-inline Tensor & dispatch_clamp_max(const Tensor & self, Scalar max, Tensor result) {
-  AutoNoGIL no_gil;
-  return at::clamp_max_out(result, self, max);
-}
-
 ${py_method_dispatch}
 
 }} // namespace torch::autograd

tools/autograd/templates/python_variable_methods.cpp

@@ -55,80 +55,6 @@ static PyObject * THPVariable_apply_(PyObject* self, PyObject* arg)
   END_HANDLE_TH_ERRORS
 }
 
-static Tensor dispatch_clamp(const Tensor & self, Scalar min, Scalar max) {
-  AutoNoGIL no_gil;
-  DeviceGuard device_guard(self);
-  return self.clamp(min, max);
-}
-static Tensor dispatch_clamp_min(const Tensor & self, Scalar min) {
-  AutoNoGIL no_gil;
-  DeviceGuard device_guard(self);
-  return self.clamp_min(min);
-}
-static Tensor dispatch_clamp_max(const Tensor & self, Scalar max) {
-  AutoNoGIL no_gil;
-  DeviceGuard device_guard(self);
-  return self.clamp_max(max);
-}
-
-static PyObject * THPVariable_clamp(PyObject* self, PyObject* args, PyObject* kwargs)
-{
-  HANDLE_TH_ERRORS
-  static PythonArgParser parser({
-    "clamp(Scalar min=None, Scalar max=None)",
-  }, /*traceable=*/true);
-  auto& self_ = reinterpret_cast<THPVariable*>(self)->cdata;
-  ParsedArgs<2> parsed_args;
-  auto r = parser.parse(args, kwargs, parsed_args);
-  if (!r.isNone(0) && !r.isNone(1)) {
-    return THPVariable_Wrap(dispatch_clamp(self_, r.scalar(0), r.scalar(1)));
-  } else if (!r.isNone(0)) {
-    return THPVariable_Wrap(dispatch_clamp_min(self_, r.scalar(0)));
-  } else if (!r.isNone(1)) {
-    return THPVariable_Wrap(dispatch_clamp_max(self_, r.scalar(1)));
-  } else {
-    throw std::runtime_error("At least one of 'min' or 'max' must not be None");
-  }
-  END_HANDLE_TH_ERRORS
-}
-
-static Tensor & dispatch_clamp_(Tensor & self, Scalar min, Scalar max) {
-  AutoNoGIL no_gil;
-  DeviceGuard device_guard(self);
-  return self.clamp_(min, max);
-}
-static Tensor & dispatch_clamp_min_(Tensor & self, Scalar min) {
-  AutoNoGIL no_gil;
-  DeviceGuard device_guard(self);
-  return self.clamp_min_(min);
-}
-static Tensor & dispatch_clamp_max_(Tensor & self, Scalar max) {
-  AutoNoGIL no_gil;
-  DeviceGuard device_guard(self);
-  return self.clamp_max_(max);
-}
-
-static PyObject * THPVariable_clamp_(PyObject* self, PyObject* args, PyObject* kwargs)
-{
-  HANDLE_TH_ERRORS
-  static PythonArgParser parser({
-    "clamp_(Scalar min=None, Scalar max=None)",
-  }, /*traceable=*/true);
-  auto& self_ = reinterpret_cast<THPVariable*>(self)->cdata;
-  ParsedArgs<2> parsed_args;
-  auto r = parser.parse(args, kwargs, parsed_args);
-  if (!r.isNone(0) && !r.isNone(1)) {
-    return THPVariable_Wrap(dispatch_clamp_(self_, r.scalar(0), r.scalar(1)));
-  } else if (!r.isNone(0)) {
-    return THPVariable_Wrap(dispatch_clamp_min_(self_, r.scalar(0)));
-  } else if (!r.isNone(1)) {
-    return THPVariable_Wrap(dispatch_clamp_max_(self_, r.scalar(1)));
-  } else {
-    throw std::runtime_error("At least one of 'min' or 'max' must not be None");
-  }
-  END_HANDLE_TH_ERRORS
-}
-
 static PyObject * THPVariable_size(PyObject* self, PyObject* args, PyObject* kwargs)
 {
   HANDLE_TH_ERRORS
@@ -661,8 +587,6 @@ PyMethodDef variable_methods[] = {
   {"apply_", (PyCFunction)THPVariable_apply_, METH_O, NULL},
   {"byte", (PyCFunction)THPVariable_byte, METH_NOARGS, NULL},
   {"char", (PyCFunction)THPVariable_char, METH_NOARGS, NULL},
-  {"clamp", (PyCFunction)THPVariable_clamp, METH_VARARGS | METH_KEYWORDS, NULL},
-  {"clamp_", (PyCFunction)THPVariable_clamp_, METH_VARARGS | METH_KEYWORDS, NULL},
   {"contiguous", (PyCFunction)THPVariable_contiguous, METH_NOARGS, NULL},
   {"copy_", (PyCFunction)THPVariable_copy_, METH_VARARGS | METH_KEYWORDS, NULL},
   {"cpu", (PyCFunction)THPVariable_cpu, METH_NOARGS, NULL},

torch/csrc/jit/interned_strings.h

@@ -19,6 +19,7 @@ _(namespaces, scope) \
 _(namespaces, namespaces) \
 _(prim, Assign) \
 _(prim, Constant) \
+_(prim, None) \
 _(prim, Drop) \
 _(prim, Eval) \
 _(prim, Expand) /* onnx */ \

torch/csrc/jit/operator.cpp

@@ -104,6 +104,9 @@ struct SchemaParser {
       case TK_FALSE:
         L.next();
         return false;
+      case TK_NONE:
+        L.next();
+        return IValue();
       case TK_IDENT: {
         auto tok = L.next();
         auto text = tok.text();
@@ -160,11 +163,8 @@ struct SchemaParser {
   }
 
   IValue parseTensorDefault(const SourceRange& range) {
-    if("None" == L.expect(TK_IDENT).text()) {
-      return at::Tensor();
-    } else {
-      throw ErrorReport(range) << "invalid tensor default value";
-    }
+    L.expect(TK_NONE);
+    return IValue();
   }
   void parseDefaultValue(Argument& arg) {
     auto range = L.cur().range;