Make test pass

Author: qihqi

torchax/test/test_image.py

@@ -1,3 +1,4 @@
+from absl.testing import parameterized
 import unittest
 from typing import Tuple
 import itertools
@@ -23,12 +24,20 @@ def upsample_jit(tensor, output_size: Tuple[int, int], align_corners: bool, anti
     return torchax.interop.jax_view(tensor)
 
 
-def test_upsampling(align_corners: bool, antialias: bool, method: str):
 
-    if method == 'bilinear':
-      if align_corners:
-        return # bilinear upsampling does not support align_corners
-    
+class TestResampling(parameterized.TestCase):
+
+  @parameterized.product(
+        antialias=[
+          True,
+          False,
+        ],
+        align_corners=[
+          False,
+          True,
+        ])
+  def test_resampling_combinations_bicubic(self, antialias, align_corners):
+    method = 'bicubic'
     input_tensor = torch.rand((1, 1, 256, 512), dtype=torch.float32)
     output_size = (128, 64)
 
@@ -44,26 +53,6 @@ def test_upsampling(align_corners: bool, antialias: bool, method: str):
 
     assert torch.allclose(upsampled_tensor, upsampled_tensor_xla, atol=1e-4, rtol=1e-5), f"{method} upsampling failed with error {abs_err.max()}"
 
-class TestResampling(unittest.TestCase):
-    def test_resampling_combinations(self):
-        methods = [
-          'bicubic',
-          'bilinear',
-        ]
-        antialias_options = [
-          True,
-          False,
-        ]
-        
-        aligncorners_options = [
-          False,
-          True,
-        ]
-
-        for method, antialias, align_corners in itertools.product(methods, antialias_options, aligncorners_options):
-            with self.subTest(method=method, antialias=antialias, align_corners=align_corners):
-                test_upsampling(align_corners=align_corners, antialias=antialias, method=method)
-        
 
 
 if __name__ == '__main__':

torchax/torchax/ops/jaten.py

@@ -5270,6 +5270,25 @@ def _aten_upsample_billinear_aa(input,
       scales_w
   )
 
+@op(torch.ops.aten._upsample_bicubic2d_aa)
+def _aten_upsample_bicubic2d_aa(input,
+                                output_size,
+                                align_corners,
+                                scale_factors=None,
+                                scales_h=None,
+                                scales_w=None):
+  return _aten_upsample(
+      input,
+      output_size,
+      align_corners,
+      True, # antialias
+      "bicubic", # method
+      scale_factors,
+      scales_h,
+      scales_w
+  )
+
+
 @op(torch.ops.aten.polar)
 def _aten_polar(abs, angle, *, out=None):
   return jax.lax.complex(abs * jnp.cos(angle), abs * jnp.sin(angle))

torchax/torchax/ops/jimage.py

@@ -13,14 +13,14 @@ def cubic_kernel(x, a=-0.75):
     return jnp.where(cond1, f1, jnp.where(cond2, f2, 0.0))
 
 
-def compute_contribs(in_size, out_size, scale, support=2.0, align_corners=False):
+def compute_contribs(in_size, out_size, scale, support=2.0, align_corners=False, dtype=None):
     if align_corners:
         if out_size == 1:
-            in_coords = jnp.zeros((1,))
+            in_coords = jnp.zeros((1,), dtype=dtype)
         else:
-            in_coords = jnp.linspace(0, in_size - 1, out_size)
+            in_coords = jnp.linspace(0, in_size - 1, out_size, dtype=dtype)
     else:
-        out_coords = jnp.arange(out_size) + 0.5
+        out_coords = jnp.arange(out_size, dtype=dtype) + 0.5
         in_coords = out_coords / scale - 0.5
 
     left_idx = jnp.floor(in_coords).astype(jnp.int32) - 1
@@ -86,11 +86,13 @@ def interpolate_bicubic_no_aa(img, out_h, out_w, align_corners=False):
 
     idxs_y, weights_y = compute_contribs(
         h, out_h, scale_y, align_corners=align_corners,
+        dtype=img.dtype,
     )
     tmp = interpolate_along_axis_bchw(img, idxs_y, weights_y, axis=2)
 
     idxs_x, weights_x = compute_contribs(
         w, out_w, scale_x, align_corners=align_corners,
+        dtype=img.dtype,
     )
     out = interpolate_along_axis_bchw(tmp, idxs_x, weights_x, axis=3)
     return out

torchax/torchax/ops/jtorch.py

@@ -545,18 +545,15 @@ def functional_interpolate(
   antialias = antialias or False
   align_corners = align_corners or False
 
-  if mode in ('cubic', 'bicubic', 'tricubic') and not antialias:
+  if mode in ('cubic', 'bicubic', 'tricubic') and not antialias and size is not None:
     return jimage.interpolate_bicubic_no_aa(
       input,
       size[0],
       size[1],
       align_corners,
     )
-  return jaten._aten_upsample(
-    input,
-    size,
-    align_corners,
-    antialias,
-    mode,
-    scale_factor,
-  )
+  else:
+    # fallback
+    raise torchax.tensor.OperatorNotFound(
+        f"JAX does not support interpolation mode: {mode}. Supported modes are: {supported_methods}"
+    )