TorchScript to support == None, != None

[dzhulgakov]

🐛 Bug

We do support is None, is not None in TorchScript. We should also allow == None and != None because unless the comparison operator is overwritten for the left hand side they behave equivalently.

Note that PEP8 recommends is version, but both work in python (https://docs.quantifiedcode.com/python-anti-patterns/readability/comparison_to_none.html)

To Reproduce

import torch
from typing import Optional

@torch.jit.script
def bar(x : Optional[torch.Tensor]):
  if x != None:
    x = torch.tensor([1,2,3])
  return x

bar(None)

produces

RuntimeError: 
Arguments for call are not valid.
The following variants are available:

  aten::ne.Tensor(Tensor self, Tensor other) -> (Tensor):
  Expected a value of type 'Tensor' for argument 'self' but instead found type 'Optional[Tensor]'.
...

Expected behavior

Should have the same behavior as is None.

cc @suo

[AvivSham]

Same behaviour here:

from typing import Optional
import torch
from torch import nn 

class m(nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self,x: torch.Tensor, idx: Optional[int]):
        if idx == 0:
            return x[idx]
        else:
            return x

torch.jit.script(m(), (torch.rand(20,20), 0))

How can we solve this?

@houseroad @suo @eellison

[ProGamerGov]

Torchscript should also support bool functionality of NoneTypes.

module: Optional[Callable] = None

if module:
   # do stuff..

Currently you get the following error:

E       RuntimeError: 
E       Could not cast value of type NoneType to bool: