Use the builtin bernoulli test | test(torchlib) (#850)

There is builtin test for the PyTorch op `bernoulli`. We should use that
to ensure consistency with PyTorch.

Also updated deterministic tests for bernoulli_p.

Author: justinchuby

onnxscript/tests/function_libs/torch_lib/extra_opinfo.py

@@ -501,7 +501,7 @@ def sample_inputs_bernoulli_p(op_info, device, dtype, requires_grad, **kwargs):
             yield opinfo_core.SampleInput(t, kwargs={"p": p})
 
 
-def sample_inputs_bernoulli_default(op_info, device, dtype, requires_grad, **kwargs):
+def sample_inputs_bernoulli_p_deterministic(op_info, device, dtype, requires_grad, **kwargs):
     del op_info
 
     shapes = [
@@ -512,16 +512,18 @@ def sample_inputs_bernoulli_default(op_info, device, dtype, requires_grad, **kwa
     ]
 
     for shape in shapes:
-        t = torch_testing.make_tensor(
-            shape,
-            low=0,
-            high=1,
-            device=device,
-            dtype=dtype,
-            requires_grad=requires_grad,
-            **kwargs,
-        )
-        yield opinfo_core.SampleInput(t)
+        for p in (0, 1):
+            t = torch_testing.make_tensor(
+                shape,
+                low=0,
+                high=1,
+                device=device,
+                dtype=dtype,
+                requires_grad=requires_grad,
+                **kwargs,
+            )
+            yield opinfo_core.SampleInput(t, args=(p,))
+            yield opinfo_core.SampleInput(t, kwargs={"p": p})
 
 
 OP_DB: List[opinfo_core.OpInfo] = [
@@ -680,10 +682,11 @@ def sample_inputs_bernoulli_default(op_info, device, dtype, requires_grad, **kwa
         sample_inputs_func=sample_inputs_bernoulli_p,
     ),
     opinfo_core.OpInfo(
-        "aten.bernoulli",
-        aten_name="bernoulli",
-        op=torch.ops.aten.bernoulli.default,
-        dtypes=common_dtype.floating_types(),
-        sample_inputs_func=sample_inputs_bernoulli_default,
+        # Deterministic bernoulli sampling where p is either 0 or 1
+        "aten.bernoulli.p_deterministic",
+        aten_name="bernoulli.p",
+        op=torch.ops.aten.bernoulli.p,
+        dtypes=common_dtype.all_types(),
+        sample_inputs_func=sample_inputs_bernoulli_p_deterministic,
     ),
 ]

onnxscript/tests/function_libs/torch_lib/ops_test_data.py

@@ -495,15 +495,7 @@ def _where_input_wrangler(
         reason="atleast_3d_single_tensor overload takes single tensor as input",
     ),
     TorchLibOpInfo("baddbmm", core_ops.aten_baddbmm),
-    TorchLibOpInfo(
-        # This string is a unique ID. In extra_opinfo.py, we
-        # also define test data for this ID with
-        # `opinfo_core.OpInfo("aten.bernoulli.p", ...)`.
-        "aten.bernoulli",
-        core_ops.aten_bernoulli,
-        # Skip comparison for the output of this op because it is a random tensor.
-        nondeterministic=True,
-    ),
+    TorchLibOpInfo("bernoulli", core_ops.aten_bernoulli, nondeterministic=True),
     TorchLibOpInfo(
         # This string is a unique ID. In extra_opinfo.py, we
         # also define test data for this ID with
@@ -513,6 +505,7 @@ def _where_input_wrangler(
         # Skip comparison for the output of this op because it is a random tensor.
         nondeterministic=True,
     ),
+    TorchLibOpInfo("aten.bernoulli.p_deterministic", core_ops.aten_bernoulli_p),
     TorchLibOpInfo("bmm", core_ops.aten_bmm),
     TorchLibOpInfo("broadcast_to", core_ops.aten_broadcast_to),
     TorchLibOpInfo("cat", core_ops.aten_cat).skip(