Fix NJT OpInfo entry for nn.functional.prelu (pytorch#144582)

Part of my BE project addressing NJT bugs surfaced via OpInfo tests.

The OpInfo entry for prelu was wrong before this PR; `weight` needs to be passed as well. The op isn't fully implemented yet.
Pull Request resolved: pytorch#144582
Approved by: https://github.com/soulitzer

Author: jbschlosser

test/test_nestedtensor.py

@@ -7909,6 +7909,7 @@ def __torch_dispatch__(self, func, types, args=..., kwargs=None):
             # unary
             # needs log_sigmoid_forward, which returns a tuple
             "nn.functional.logsigmoid",
+            "nn.functional.prelu",
             # needs rrelu_with_noise
             "nn.functional.rrelu",
             # binary
@@ -7972,13 +7973,6 @@ def __torch_dispatch__(self, func, types, args=..., kwargs=None):
         },
         name="no_masked_jagged_support",
     ),
-    # Need to adjust sample input func to pass the right thing
-    XFailRule(
-        error_type=TypeError,
-        error_msg="missing 1 required positional arguments",
-        op_match_fn=lambda device, op: (op.full_name == "nn.functional.prelu"),
-        name="invalid_prelu_sample_input_func",
-    ),
     # Op doesn't support lengths being present
     XFailRule(
         error_type=ValueError,

torch/testing/_internal/opinfo/definitions/nested.py

@@ -1217,6 +1217,41 @@ def sample_inputs_nn_functional_linear(op_info, device, dtype, requires_grad, **
         )
 
 
+def sample_inputs_nn_functional_prelu(op_info, device, dtype, requires_grad, **kwargs):
+    for njt in _sample_njts(
+        device=device, dtype=dtype, requires_grad=requires_grad, dims=[3, 4]
+    ):
+        # Second dim is interpreted as number of channels; this should be non-ragged for now
+        num_channels = njt.size(1)
+        if is_nested_int(num_channels):
+            continue
+
+        # 1D weight
+        weight = torch.randn(
+            num_channels,
+            device=device,
+            dtype=dtype,
+            requires_grad=requires_grad,
+        )
+
+        yield SampleInput(
+            _clone(njt),
+            kwargs={
+                "weight": _clone(weight),
+            },
+            name=f"{_describe_njt(njt)}: 1D weight",
+        )
+
+        # scalar tensor weight
+        yield SampleInput(
+            _clone(njt),
+            kwargs={
+                "weight": torch.tensor(4.2, device=device, dtype=dtype),
+            },
+            name=f"{_describe_njt(njt)}: scalar tensor weight",
+        )
+
+
 def sample_inputs_nn_functional_rms_norm(
     op_info, device, dtype, requires_grad, **kwargs
 ):
@@ -1412,6 +1447,7 @@ def sample_inputs_where(op_info, device, dtype, requires_grad, **kwargs):
     "nn.functional.embedding": sample_inputs_nn_functional_embedding,
     "nn.functional.embedding_bag": sample_inputs_nn_functional_embedding_bag,
     "nn.functional.linear": sample_inputs_nn_functional_linear,
+    "nn.functional.prelu": sample_inputs_nn_functional_prelu,
     "nn.functional.rms_norm": sample_inputs_nn_functional_rms_norm,
     "nn.functional.threshold": sample_inputs_nn_functional_threshold,
     **{f"polygamma.polygamma_n_{n}": sample_inputs_polygamma_n(n=n) for n in range(5)},