adding mode_list argument in _interp_args so custom list of modes…

… can be drawn for resampling frontend functions
IgorTatarnikov · Mar 8, 2023 · c8dd1c5 · c8dd1c5
1 parent 3aed083
commit c8dd1c5
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/ivy_tests/test_ivy/test_frontends/test_torch/test_vision_functions.py b/ivy_tests/test_ivy/test_frontends/test_torch/test_vision_functions.py
@@ -160,7 +160,10 @@ def test_torch_pad(
 
 @handle_frontend_test(
     fn_tree="torch.nn.functional.interpolate",
-    dtype_and_input_and_other=_interp_args(scale_factor=True),
+    dtype_and_input_and_other=_interp_args(
+        mode_list=["linear", "bilinear", "trilinear", "nearest", "area"],
+        scale_factor=True,
+    ),
     number_positional_args=st.just(2),
 )
 def test_torch_interpolate(

diff --git a/ivy_tests/test_ivy/test_functional/test_experimental/test_nn/test_layers.py b/ivy_tests/test_ivy/test_functional/test_experimental/test_nn/test_layers.py
@@ -257,13 +257,15 @@ def test_dct(
 
 
 @st.composite
-def _interp_args(draw, mode=None, scale_factor=False):
-    if not mode:
+def _interp_args(draw, mode=None, mode_list=None, scale_factor=False):
+    if not mode and not mode_list:
         mode = draw(
             st.sampled_from(
                 ["linear", "bilinear", "trilinear", "nearest", "area", "tf_area"]
             )
         )
+    elif mode_list:
+        mode = draw(st.sampled_from(mode_list))
     align_corners = draw(st.one_of(st.booleans(), st.none()))
     if mode == "linear":
         num_dims = 3