Fix train dtype for native_dropout | fix(torchlib) (#869)

Previously the `train` attribute in native_dropout was not casted to
BOOL. We need to do that because the underlying attribute type for
python bool is INT64 when promoted to input. This change also added a
test that will catch the old error.

onnxscript/function_libs/torch_lib/ops/core.py

@@ -4609,6 +4609,10 @@ def aten_native_dropout(
 ) -> Tuple[TFloatOrBFloat16, BOOL]:
     """native_dropout(Tensor input, float p, bool? train) -> (Tensor, Tensor)"""
 
+    # Python bool attributes need to be explicitly converted to BOOL
+    # because the underlying attribute type is int
+    # TODO(#872): Allow ONNX Script to handle this conversion
+    train = op.Cast(train, to=BOOL.dtype)
     result, mask = op.Dropout(input, p, train)
     return result, mask
 

onnxscript/tests/function_libs/torch_lib/extra_opinfo.py

@@ -16,6 +16,8 @@
 )
 from torch.testing._internal.opinfo import core as opinfo_core
 
+S = 5
+
 
 def sample_inputs__local_scalar_dense(op_info, device, dtype, requires_grad, **kwargs):
     del op_info
@@ -446,6 +448,28 @@ def sample_inputs_col2im(op_info, device, dtype, requires_grad, **kwargs):
         yield opinfo_core.SampleInput(tensor, args=(output_size, kernel_size), kwargs=kwargs)
 
 
+def sample_inputs_native_dropout(
+    op_info, device, dtype, requires_grad, *, valid_input_dim=None, **kwargs
+):
+    del op_info  # Unused
+    del kwargs  # Unused
+    make_arg = functools.partial(
+        torch_testing.make_tensor, device=device, dtype=dtype, requires_grad=requires_grad
+    )
+
+    if valid_input_dim:
+        cases = ((S,) * i for i in valid_input_dim)
+    else:
+        cases = ((S, S), (S,), ())
+    # ONNX requires 0 <= p < 1
+    p_vals = [0.0]
+
+    training_vals = [True, False]
+
+    for case, p, training in itertools.product(cases, p_vals, training_vals):
+        yield opinfo_core.SampleInput(make_arg(case), p=p, train=training)
+
+
 def sample_inputs_stft(op_info, device, dtype, requires_grad, **kwargs):
     del op_info
     del kwargs
@@ -561,6 +585,15 @@ def sample_inputs_bernoulli_p_deterministic(op_info, device, dtype, requires_gra
         dtypes=common_dtype.all_types(),
         sample_inputs_func=sample_inputs__local_scalar_dense,
     ),
+    opinfo_core.OpInfo(
+        "ops.aten.native_dropout",
+        aten_name="native_dropout",
+        dtypes=common_dtype.all_types_and_half(),
+        sample_inputs_func=sample_inputs_native_dropout,
+        supports_out=False,
+    ),
+    # TODO(#870): Rename all with the op.aten. prefix so OpInfo will find
+    # the op automatically.
     opinfo_core.OpInfo(
         "col2im",
         op=torch.ops.aten.col2im,

onnxscript/tests/function_libs/torch_lib/ops_test_data.py

@@ -748,7 +748,7 @@ def _where_input_wrangler(
     TorchLibOpInfo("mm", core_ops.aten_mm),
     TorchLibOpInfo("mul", core_ops.aten_mul),
     TorchLibOpInfo("narrow", core_ops.aten_narrow),
-    # TorchLibOpInfo("native_dropout", core_ops.aten_native_dropout),  # native_dropout is not in OPS_DB
+    TorchLibOpInfo("ops.aten.native_dropout", core_ops.aten_native_dropout),
     TorchLibOpInfo("ne", core_ops.aten_ne),
     TorchLibOpInfo("neg", core_ops.aten_neg),
     TorchLibOpInfo(