diff --git a/onnxscript/function_libs/torch_lib/ops/nn.py b/onnxscript/function_libs/torch_lib/ops/nn.py index 0b8d6cba50..8bd0e42699 100644 --- a/onnxscript/function_libs/torch_lib/ops/nn.py +++ b/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, diff --git a/tests/function_libs/torch_lib/extra_opinfo.py b/tests/function_libs/torch_lib/extra_opinfo.py index 5f4b9aa486..0d01800a24 100644 --- a/tests/function_libs/torch_lib/extra_opinfo.py +++ b/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", diff --git a/tests/function_libs/torch_lib/ops_test_data.py b/tests/function_libs/torch_lib/ops_test_data.py index 15a0cec84c..b08874579a 100644 --- a/tests/function_libs/torch_lib/ops_test_data.py +++ b/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),