[torchlib] Add aten::upsample_linear1d.vec

onnxscript/function_libs/torch_lib/ops/nn.py

@@ -2588,6 +2588,32 @@ def aten_upsample_linear1d(
     )
 
 
+@torch_op("aten::upsample_linear1d.vec", trace_only=True)
+def aten_upsample_linear1d_vec(
+    self: TReal,
+    output_size: Optional[INT64],
+    align_corners: bool,
+    scale_factors: Optional[Sequence[float]],
+) -> TReal:
+    """upsample_linear1d.vec(Tensor input, SymInt[]? output_size, bool align_corners, float[]? scale_factors) -> Tensor"""
+
+    coordinate_transformation_mode = _get_upsample_align_corners_mode(align_corners)
+    if scale_factors is not None:
+        return _aten_upsample_scales(
+            self,
+            scale_factors,
+            mode="linear",
+            coordinate_transformation_mode=coordinate_transformation_mode,
+        )
+    assert output_size is not None
+    return _aten_upsample_output_size(
+        self,
+        output_size,
+        mode="linear",
+        coordinate_transformation_mode=coordinate_transformation_mode,
+    )
+
+
 def aten_upsample_linear1d_backward(
     grad_output: TensorType,
     output_size: INT64,

tests/function_libs/torch_lib/extra_opinfo.py

@@ -2086,6 +2086,49 @@ def shape(size, rank, with_batch_channel=True):
         )
 
 
+def sample_inputs_upsample_linear1d_vec(op_info, device, dtype, requires_grad, **kwargs):
+    del op_info
+    del kwargs
+
+    N, C = 2, 3
+    D = 4
+    SS = 3
+    L = 5
+
+    align_corners_options = (True, False)
+    rank = 1
+
+    def shape(size, rank, with_batch_channel=True):
+        if with_batch_channel:
+            return tuple([N, C] + ([size] * rank))
+        return tuple([size] * rank)
+
+    make_arg = functools.partial(
+        torch_testing.make_tensor,
+        device=device,
+        dtype=dtype,
+        requires_grad=requires_grad,
+        low=-1,
+        high=1,
+    )
+
+    yield opinfo_core.SampleInput(make_arg(shape(D, rank)), shape(SS, rank, False), True, None)
+
+    for align_corners in align_corners_options:
+        yield opinfo_core.SampleInput(
+            make_arg(shape(D, rank)), shape(1, rank, False), align_corners, None
+        )
+        yield opinfo_core.SampleInput(
+            make_arg(shape(D, rank)), shape(L, rank, False), align_corners, None
+        )
+        for scale_factors in ([1.7], [0.6], [4.2]):
+            yield opinfo_core.SampleInput(
+                make_arg(shape(D, rank)),
+                args=(None, align_corners),  # output_size must be None with scale_factors
+                kwargs=dict(scale_factors=scale_factors),
+            )
+
+
 def sample_inputs_upsample_nearest1d(op_info, device, dtype, requires_grad, **kwargs):
     del op_info
     del kwargs
@@ -3099,6 +3142,13 @@ def sample_inputs_masked_scatter(op_info, device, dtype, requires_grad, **kwargs
         sample_inputs_func=sample_inputs_upsample_linear1d,
         supports_out=False,
     ),
+    opinfo_core.OpInfo(
+        "ops.aten.upsample_linear1d.vec",
+        aten_name="upsample_linear1d.vec",
+        dtypes=common_dtype.floating_types_and(torch.bfloat16),
+        sample_inputs_func=sample_inputs_upsample_linear1d_vec,
+        supports_out=False,
+    ),
     opinfo_core.OpInfo(
         "ops.aten.upsample_nearest1d",
         aten_name="upsample_nearest1d",

tests/function_libs/torch_lib/ops_test_data.py

@@ -1830,6 +1830,7 @@ def _where_input_wrangler(
         ),
         reason="fixme: align_corners=False output mismatch when scales are provided",
     ),
+    TorchLibOpInfo("ops.aten.upsample_linear1d.vec", nn_ops.aten_upsample_linear1d_vec),
     TorchLibOpInfo("ops.aten.upsample_nearest1d", nn_ops.aten_upsample_nearest1d),
     TorchLibOpInfo("ops.aten.upsample_nearest1d.vec", nn_ops.aten_upsample_nearestnd_vec),
     TorchLibOpInfo("ops.aten.upsample_nearest2d", nn_ops.aten_upsample_nearest2d),