Fix comparison functions

Author: apaszke

test/test_autograd.py

@@ -682,6 +682,24 @@ def test_inplace(self):
  x.add_(2)
  self.assertRaises(RuntimeError, lambda: z.backward(torch.ones(5, 5)))
 
+ def test_mark_non_differentiable(self):
+ class MyFunction(Function):
+ @staticmethod
+ def forward(ctx, input):
+ output = input > 0
+ ctx.mark_non_differentiable(output)
+ return output
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ return (grad_output * 0).type(torch.DoubleTensor)
+
+ x = Variable(torch.randn(5, 5), requires_grad=True)
+ mask = MyFunction.apply(x)
+ self.assertFalse(mask.requires_grad)
+ y = x.masked_fill(mask, 0)
+ y.sum().backward()
+
  def test_shared_storage(self):
  x = Variable(torch.ones(5, 5))
  y = x.t()

torch/autograd/_functions/compare.py

@@ -3,18 +3,22 @@
 from ..function import Function
 
 
+# TODO: once Cpp-style functions are implemented we can detach a and b
+# before calling forward.
 class _CompareOp(Function):
 
- def __init__(self, scalar=None):
- super(_CompareOp, self).__init__()
- self.scalar = scalar
-
- def forward(self, tensor1, tensor2=None):
- other = tensor2 if tensor2 is not None else self.scalar
- mask = getattr(tensor1, self.fn_name)(other)
- self.mark_non_differentiable(mask)
+ @classmethod
+ def forward(cls, ctx, a, b):
+ ctx.b_tensor = torch.is_tensor(b)
+ mask = getattr(a, cls.fn_name)(b)
+ ctx.mark_non_differentiable(mask)
  return mask
 
+ @staticmethod
+ def backward(ctx, grad_output):
+ grad_input = grad_output * 0
+ return grad_input, (grad_input if ctx.b_tensor else None)
+
 
 class Eq(_CompareOp):
  fn_name = 'eq'

torch/autograd/function.py

@@ -94,8 +94,9 @@ class FunctionMeta(type):
 
  def __init__(cls, name, bases, attrs):
  for super_cls in cls.mro():
- if 'forward' in super_cls.__dict__:
- has_static_forward = isinstance(super_cls.__dict__['forward'], staticmethod)
+ forward = super_cls.__dict__.get('forward')
+ if forward is not None:
+ has_static_forward = isinstance(forward, staticmethod) or isinstance(forward, classmethod)
  break
 
  setattr(cls, '_is_legacy', not has_static_forward)

torch/autograd/variable.py

@@ -706,40 +706,28 @@ def bernoulli(self):
  return Bernoulli()(self)
 
  def eq(self, other):
- if isinstance(other, Variable):
- return Eq()(self, other)
  assert not torch.is_tensor(other), "can't compare Variable and tensor"
- return Eq(other)(self)
+ return Eq.apply(self, other)
 
  def ne(self, other):
- if isinstance(other, Variable):
- return Ne()(self, other)
  assert not torch.is_tensor(other), "can't compare Variable and tensor"
- return Ne(other)(self)
+ return Ne.apply(self, other)
 
  def gt(self, other):
- if isinstance(other, Variable):
- return Gt()(self, other)
  assert not torch.is_tensor(other), "can't compare Variable and tensor"
- return Gt(other)(self)
+ return Gt.apply(self, other)
 
  def ge(self, other):
- if isinstance(other, Variable):
- return Ge()(self, other)
  assert not torch.is_tensor(other), "can't compare Variable and tensor"
- return Ge(other)(self)
+ return Ge.apply(self, other)
 
  def lt(self, other):
- if isinstance(other, Variable):
- return Lt()(self, other)
  assert not torch.is_tensor(other), "can't compare Variable and tensor"
- return Lt(other)(self)
+ return Lt.apply(self, other)
 
  def le(self, other):
- if isinstance(other, Variable):
- return Le()(self, other)
  assert not torch.is_tensor(other), "can't compare Variable and tensor"
- return Le(other)(self)
+ return Le.apply(self, other)
 
  def __add__(self, other):
  return self.add(other)