import torch
from torch.autograd import Function, gradcheck

class Sum1(Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x.sum()

    @staticmethod
    def backward(ctx, grad_output):
        grad_x = None

        x = ctx.saved_tensors

        if ctx.needs_input_grad[0]:
            grad_x = torch.ones_like(x) * grad_output

        return grad_x

class Sum2(Function):
    @staticmethod
    def forward(ctx, x, y):
        ctx.save_for_backward(x)
        return x.sum()

    @staticmethod
    def backward(ctx, grad_output):
        grad_x = grad_y = None

        x = ctx.saved_tensors

        if ctx.needs_input_grad[0]:
            grad_x = torch.ones_like(x) * grad_output

        return grad_x, grad_y

class Sum3(Function):
    @staticmethod
    def forward(ctx, x, y):
        ctx.save_for_backward(x)
        return x.sum()

    @staticmethod
    def backward(ctx, grad_output):
        grad_x = grad_y = None

        x = ctx.saved_tensors[0]

        if ctx.needs_input_grad[0]:
            grad_x = torch.ones_like(x) * grad_output

        return grad_x, grad_y

sum1 = Sum1.apply
sum2 = Sum2.apply
sum3 = Sum3.apply

a = torch.rand(10, dtype=torch.double, requires_grad=True)
b = torch.rand(10, dtype=torch.double)

gradcheck(sum3, (a,b))
gradcheck(sum2, (a,b))
gradcheck(sum1, a)
gradcheck(sum1, (a))

~/miniconda3/envs/pytorch-nightly/lib/python3.6/site-packages/torch/autograd/gradcheck.py in gradcheck(func, inputs, eps, atol, rtol, raise_exception)
    193             return _as_tuple(func(*input))[i]
    194
--> 195         analytical, reentrant, correct_grad_sizes = get_analytical_jacobian(tupled_inputs, o)
    196         numerical = get_numerical_jacobian(fn, inputs, eps=eps)
    197

~/miniconda3/envs/pytorch-nightly/lib/python3.6/site-packages/torch/autograd/gradcheck.py in get_analytical_jacobian(input, output)
     94         for jacobian_c in (jacobian, jacobian_reentrant):
     95             grads_input = torch.autograd.grad(output, diff_input_list, grad_output,
---> 96                                               retain_graph=True, allow_unused=True)
     97             for jacobian_x, d_x, x in zip(jacobian_c, grads_input, diff_input_list):
     98                 if d_x is not None and d_x.size() != x.size():

~/miniconda3/envs/pytorch-nightly/lib/python3.6/site-packages/torch/autograd/__init__.py in grad(outputs, inputs, grad_outputs, retain_graph, create_graph, only_inputs, allow_unused)
    143     return Variable._execution_engine.run_backward(
    144         outputs, grad_outputs, retain_graph, create_graph,
--> 145         inputs, allow_unused)
    146
    147

~/miniconda3/envs/pytorch-nightly/lib/python3.6/site-packages/torch/autograd/function.py in apply(self, *args)
     74
     75     def apply(self, *args):
---> 76         return self._forward_cls.backward(self, *args)
     77
     78

<ipython-input-1-1f2fe1d5eb24> in backward(ctx, grad_output)
     32
     33         if ctx.needs_input_grad[0]:
---> 34             grad_x = torch.ones_like(x) * grad_output
     35
     36         return grad_x, grad_y

TypeError: ones_like() received an invalid combination of arguments - got (tuple), but expected one of:
 * (Tensor input, torch.dtype dtype, torch.layout layout, torch.device device, bool requires_grad)
 * (Tensor input, bool requires_grad)

~/miniconda3/envs/pytorch-nightly/lib/python3.6/site-packages/torch/autograd/gradcheck.py in gradcheck(func, inputs, eps, atol, rtol, raise_exception)
    179             'but none of the them have requires_grad=True.')
    180
--> 181     output = _differentiable_outputs(func(*inputs))
    182
    183     def fail_test(msg):

TypeError: forward() takes 2 positional arguments but 11 were given