RPC couldn't match torch.ones with requires_grad=True

[mrshenli]

🐛 Bug

I was trying to use rpc.remote to create a tensor with requires_grad=True, but hit the following error:

Couldn't match schema: aten::ones(int[] size, *, int? dtype=None, int? layout=None, Device? device=None, bool? pin_memory=None) -> (Tensor) to args: ((3, 3),) and kwar
gs: {'requires_grad': True}, reason: Unknown keyword argument 'requires_grad' for operator 'aten::ones'. Schema: aten::ones(int[] size, *, int? dtype=None, int? layout
=None, Device? device=None, bool? pin_memory=None) -> (Tensor)

To Reproduce

rpc.remote(dst, torch.ones, args=((3, 3),), kwargs={"requires_grad":True})

Work Around

Wrap torch.ones into a user function can pass that user function to remote, e.g.,

def _torch_ones(sizes, requires_grad=False):
    return torch.ones(sizes, requires_grad=requires_grad)

rpc.remote(dst, _torch_ones, args=((2, 2), ), kwargs={"requires_grad":True})

cc @ezyang @gchanan @zou3519 @bdhirsh @bhosmer @smessmer @ljk53 @ailzhang @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @rohan-varma @jjlilley @osalpekar @jiayisuse @agolynski @SciPioneer @H-Huang @mrzzd @xush6528

[rohan-varma]

Even though there's a workaround (that uses UDF over builtin types basically), this seems like its pretty important since we should be able to send over tensors that require grad for distributed autograd etc.What might be the right way to get this fixed? Should we add the optional arg to the shapes in torch/csrc/jit/passes/shape_analysis.cpp or is there another better approach?

cc JIT folks @zdevito @jamesr66a

[ezyang]

Tagging internals on this, this is hitting a longstanding problem #19393