[torchlib] Implement upsample_nearest{nd}.vec

onnxscript/function_libs/torch_lib/ops/nn.py

@@ -2355,17 +2355,6 @@
     return "align_corners" if align_corners else "pytorch_half_pixel"
 
 
-@torch_op(
-    (
-        "aten::upsample_bicubic2d",
-        "aten::upsample_bilinear2d",
-        "aten::upsample_nearest1d",
-        "aten::upsample_nearest2d",
-        "aten::upsample_nearest3d",
-        "aten::upsample_trilinear3d",
-    ),
-    private=True,
-)
 def _aten_upsample_output_size(
     self: TReal,
     output_size: INT64,
@@ -2388,7 +2377,6 @@
     )
 
 
-@torch_op(("aten::upsample_bicubic2d", "aten::upsample_bilinear2d"), private=True)
 def _aten_upsample_scales(
     self: TReal,
     scale_factors: TFloat,
@@ -2404,6 +2392,7 @@
         None,
         mode=mode,
         coordinate_transformation_mode=coordinate_transformation_mode,
+        nearest_mode="floor",
     )
 
 
@@ -2567,25 +2556,32 @@
     if size is not None:
         return _aten_upsample_output_size(self, size, "nearest", "asymmetric")
     else:
-        return _aten_upsample_nearest1d_scales(self, scale_factor)
+        return _aten_upsample_scales(
+            self, op.Constant(value_floats=[scale_factor]), "nearest", "asymmetric"
+        )
 
 
-@torch_op("aten::upsample_nearest1d", private=True)
-def _aten_upsample_nearest1d_scales(
-    self: TReal,
-    scale_factors: TFloat,
+@torch_op(
+    (
+        "aten::upsample_nearest1d.vec",
+        "aten::upsample_nearest2d.vec",
+        "aten::upsample_nearest3d.vec",
+    ),
+    trace_only=True,
+)
+def aten_upsample_nearestnd_vec(
+    input: TReal,
+    output_size: Optional[INT64] = None,
+    scale_factors: Optional[Sequence[float]] = None,
 ) -> TReal:
-    scale_factors = op.Cast(scale_factors, to=FLOAT.dtype)
-    scale_factors = op.Concat(op.Constant(value_floats=[1.0, 1.0]), scale_factors, axis=0)
-    return op.Resize(
-        self,
-        None,
-        scale_factors,  # format should be: [1.0, 1.0, scale_h, scale_w]
-        None,
-        mode="nearest",
-        coordinate_transformation_mode="asymmetric",
-        nearest_mode="floor",
-    )
+    """upsample_nearest3d.vec(Tensor input, SymInt[]? output_size, float[]? scale_factors) -> Tensor"""
+
+    if scale_factors is not None:
+        return _aten_upsample_scales(
+            input, op.Constant(value_floats=scale_factors), "nearest", "asymmetric"
+        )
+    else:
+        return _aten_upsample_output_size(input, output_size, "nearest", "asymmetric")
 
 
 def aten_upsample_nearest1d_backward(
@@ -2602,18 +2598,21 @@
 @torch_op("aten::upsample_nearest2d", trace_only=True)
 def aten_upsample_nearest2d(
     self: TReal,
-    size: INT64,
+    output_size: INT64,
     scales_h: Optional[float] = None,
     scales_w: Optional[float] = None,
 ) -> TReal:
     """upsample_nearest2d(Tensor self, SymInt[2] output_size, float? scales_h=None, float? scales_w=None) -> Tensor"""
 
-    # NOTE: trace_only because optional attributes are not supported by ONNX
-    # TODO(justinchuby): Conditionally use scales
-    del scales_h
-    del scales_w
-
-    return _aten_upsample_output_size(self, size, "nearest", "asymmetric")
+    if scales_h is not None and scales_w is not None:
+        return _aten_upsample_scales(
+            self,
+            op.Constant(value_floats=[scales_h, scales_w]),
+            "nearest",
+            "asymmetric",
+        )
+    else:
+        return _aten_upsample_output_size(self, size, "nearest", "asymmetric")
 
 
 def aten_upsample_nearest2d_backward(
@@ -2638,11 +2637,15 @@
 ) -> TReal:
     """upsample_nearest3d(Tensor self, SymInt[3] output_size, float? scales_d=None, float? scales_h=None, float? scales_w=None) -> Tensor"""
 
-    del scales_h
-    del scales_w
-    del scales_d
-
-    return _aten_upsample_output_size(self, size, "nearest", "asymmetric")
+    if scales_d is not None and scales_h is not None and scales_w is not None:
+        return _aten_upsample_scales(
+            self,
+            op.Constant(value_floats=[scales_d, scales_h, scales_w]),
+            "nearest",
+            "asymmetric",
+        )
+    else:
+        return _aten_upsample_output_size(self, size, "nearest", "asymmetric")
 
 
 def aten_upsample_nearest3d_backward(

tests/function_libs/torch_lib/ops_test_data.py

@@ -2096,14 +2096,26 @@ def _where_input_wrangler(
         "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,
     ),
+    TorchLibOpInfo(
+        "ops.aten.upsample_nearest2d.vec",
+        nn_ops.aten_upsample_nearestnd_vec,
+    ),
     TorchLibOpInfo(
         "ops.aten.upsample_nearest3d",
         nn_ops.aten_upsample_nearest3d,
     ),
+    TorchLibOpInfo(
+        "ops.aten.upsample_nearest3d.vec",
+        nn_ops.aten_upsample_nearestnd_vec,
+    ),
     TorchLibOpInfo(
         "ops.aten.upsample_trilinear3d.default",
         nn_ops.aten_upsample_trilinear3d,