First bug filed  https://github.com/pytorch/pytorch/issues/22124 11
For the case with 1-mask and with @ngimel ’ s solution:
graph(%input : Float(*, *, *, *),
      %p : float,
      %2 : bool):
  %5 : int = prim::Constant[value=0]() # foo.py:13:30
  %4 : float = prim::Constant[value=1]() # foo.py:11:19
  %3 : Scalar? = prim::Constant()
  %6 : Tensor = prim::If(%2) # foo.py:10:9
    block0():
      %p1m : float = aten::sub(%4, %p) # foo.py:11:19
      %8 : Float(*, *, *, *) = aten::rand_like(%input) # foo.py:12:20
      %mask.1 : Byte(*, *, *, *) = aten::lt(%8, %p1m) # foo.py:12:20
      %10 : Byte(*, *, *, *) = aten::gt(%input, %5) # foo.py:13:22
      %mask : Tensor = aten::mul_(%mask.1, %10) # foo.py:13:13
      %14 : float = aten::div(%4, %p1m) # foo.py:14:51
      %17 : Tensor = prim::FusionGroup_0(%input, %mask, %14)
      -> (%17)
    block1():
      %16 : Tensor = aten::clamp_(%input, %5, %3) # foo.py:16:20
      -> (%16)
  return (%6)
with prim::FusionGroup_0 = graph(%3 : Float(*, *, *, *),
      %6 : Tensor,
      %1 : float):
  %7 : Tensor = aten::type_as(%6, %3) # foo.py:14:28
  %5 : Tensor = aten::mul(%3, %7) # foo.py:14:20
  %2 : Tensor = aten::mul(%5, %1) # foo.py:14:20
  return (%2)

So if inplace gen is supported in future, it’s pretty nice.
Also impossiblity of passing a constant parameter (p) to a constructor of torch.jit.ScriptModule's derived class is quite strange (it suggests adding to __constants__, but it’s not really a “constant”).