update test convention (#14)

ivy-llc · Aug 3, 2023 · 189263c · 189263c
1 parent 3bc6dc0
commit 189263c
Show file tree

Hide file tree

Showing 25 changed files with 278 additions and 335 deletions.
diff --git a/conftest.py b/conftest.py
@@ -1,31 +1,23 @@
 # global
 import pytest
-from typing import Dict
 
 # local
 import ivy
-from ivy_tests.test_ivy import helpers
+import jax
 
-
-FW_STRS = ["numpy", "jax", "tensorflow", "torch"]
+jax.config.update("jax_enable_x64", True)
 
 
-TEST_BACKENDS: Dict[str, callable] = {
-    "numpy": lambda: helpers.globals._get_ivy_numpy(),
-    "jax": lambda: helpers.globals._get_ivy_jax(),
-    "tensorflow": lambda: helpers.globals._get_ivy_tensorflow(),
-    "torch": lambda: helpers.globals._get_ivy_torch(),
-}
+FW_STRS = ["numpy", "jax", "tensorflow", "torch"]
 
 
 @pytest.fixture(autouse=True)
-def run_around_tests(dev_str, f, compile_graph, implicit, fw):
-    if "gpu" in dev_str and fw == "numpy":
+def run_around_tests(device, compile_graph, fw):
+    if "gpu" in device and fw == "numpy":
         # Numpy does not support GPU
         pytest.skip()
-    ivy.unset_backend()
-    with f.use:
-        with ivy.DefaultDevice(dev_str):
+    with ivy.utils.backend.ContextManager(fw):
+        with ivy.DefaultDevice(device):
             yield
 
 
@@ -40,7 +32,7 @@ def pytest_generate_tests(metafunc):
     # framework
     raw_value = metafunc.config.getoption("--backend")
     if raw_value == "all":
-        backend_strs = TEST_BACKENDS.keys()
+        backend_strs = FW_STRS
     else:
         backend_strs = raw_value.split(",")
 
@@ -53,33 +45,16 @@ def pytest_generate_tests(metafunc):
     else:
         compile_modes = [False]
 
-    # with_implicit
-    raw_value = metafunc.config.getoption("--with_implicit")
-    if raw_value == "true":
-        implicit_modes = [True, False]
-    else:
-        implicit_modes = [False]
-
     # create test configs
     configs = list()
     for backend_str in backend_strs:
         for device in devices:
             for compile_graph in compile_modes:
-                for implicit in implicit_modes:
-                    configs.append(
-                        (
-                            device,
-                            TEST_BACKENDS[backend_str](),
-                            compile_graph,
-                            implicit,
-                            backend_str,
-                        )
-                    )
-    metafunc.parametrize("dev_str,f,compile_graph,implicit,fw", configs)
+                configs.append((device, compile_graph, backend_str))
+    metafunc.parametrize("device,compile_graph,fw", configs)
 
 
 def pytest_addoption(parser):
     parser.addoption("--device", action="store", default="cpu")
     parser.addoption("--backend", action="store", default="numpy,jax,tensorflow,torch")
     parser.addoption("--compile_graph", action="store", default="true")
-    parser.addoption("--with_implicit", action="store", default="false")
diff --git a/ivy_vision/implicit.py b/ivy_vision/implicit.py
@@ -52,7 +52,7 @@ def create_sampled_pixel_coords_image(
     normalized=False,
     randomize=True,
     homogeneous=False,
-    dev_str=None,
+    device=None,
 ):
     r"""Create image of randomly sampled homogeneous integer :math:`xy` pixel
     co-ordinates :math:`\mathbf{X}\in\mathbb{Z}^{h×w×3}`, stored as floating point
@@ -81,7 +81,7 @@ def create_sampled_pixel_coords_image(
     homogeneous
         Whether the pixel co-ordinates should be 3D homogeneous or just 2D.
         Default is True.
-    dev_str
+    device
         device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc.
         (Default value = None)
 
@@ -93,7 +93,7 @@ def create_sampled_pixel_coords_image(
     """
     # BS x DH x DW x 2
     low_res_pix_coords = ivy_svg.create_uniform_pixel_coords_image(
-        samples_per_dim, batch_shape, homogeneous=False, dev_str=dev_str
+        samples_per_dim, batch_shape, homogeneous=False, device=device
     )
 
     # 2
@@ -106,7 +106,7 @@ def create_sampled_pixel_coords_image(
                 ]
             )
         ),
-        device=dev_str,
+        device=device,
     )
 
     # BS x DH x DW x 2
@@ -118,27 +118,27 @@ def create_sampled_pixel_coords_image(
             low=ivy.to_scalar(-window_size[0] / 2),
             high=ivy.to_scalar(window_size[0] / 2),
             shape=list(downsam_pix_coords.shape[:-1]) + [1],
-            device=dev_str,
+            device=device,
         )
         rand_y = ivy.random_uniform(
             low=ivy.to_scalar(-window_size[1] / 2),
             high=ivy.to_scalar(window_size[1] / 2),
             shape=list(downsam_pix_coords.shape[:-1]) + [1],
-            device=dev_str,
+            device=device,
         )
 
         # BS x DH x DW x 2
         rand_offsets = ivy.concat((rand_x, rand_y), axis=-1)
         downsam_pix_coords += rand_offsets
     downsam_pix_coords = ivy.clip(
         ivy.round(downsam_pix_coords),
-        ivy.array([0.0] * 2, device=dev_str),
-        ivy.array(list(reversed(image_dims)), dtype="float32", device=dev_str) - 1,
+        ivy.array([0.0] * 2, device=device),
+        ivy.array(list(reversed(image_dims)), dtype="float32", device=device) - 1,
     )
 
     if normalized:
         downsam_pix_coords /= (
-            ivy.array([image_dims[1], image_dims[0]], dtype="float32", device=dev_str)
+            ivy.array([image_dims[1], image_dims[0]], dtype="float32", device=device)
             + MIN_DENOMINATOR
         )
 
@@ -152,7 +152,7 @@ def create_sampled_pixel_coords_image(
     return downsam_pix_coords
 
 
-def sample_images(list_of_images, num_pixels, batch_shape, image_dims, dev_str=None):
+def sample_images(list_of_images, num_pixels, batch_shape, image_dims, device=None):
     """Samples each image in a list of aligned images at num_pixels random pixel
     co-ordinates, within a unifrom grid over the image.
 
@@ -166,7 +166,7 @@ def sample_images(list_of_images, num_pixels, batch_shape, image_dims, dev_str=N
         Shape of batch. Inferred from inputs if None.
     image_dims
         Image dimensions. Inferred from inputs in None.
-    dev_str
+    device
         device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc.
         Same as images if None. (Default value = None)
 
@@ -177,8 +177,8 @@ def sample_images(list_of_images, num_pixels, batch_shape, image_dims, dev_str=N
     if image_dims is None:
         image_dims = list_of_images[0].shape[-3:-1]
 
-    if dev_str is None:
-        dev_str = ivy.dev(list_of_images[0])
+    if device is None:
+        device = ivy.dev(list_of_images[0])
 
     image_channels = [img.shape[-1] for img in list_of_images]
 
@@ -194,7 +194,7 @@ def sample_images(list_of_images, num_pixels, batch_shape, image_dims, dev_str=N
     # BS x DH x DW x 2
     sampled_pix_coords = ivy.astype(
         create_sampled_pixel_coords_image(
-            image_dims, new_img_dims, batch_shape, homogeneous=False, dev_str=dev_str
+            image_dims, new_img_dims, batch_shape, homogeneous=False, device=device
         ),
         "int32",
     )
@@ -209,8 +209,8 @@ def sample_images(list_of_images, num_pixels, batch_shape, image_dims, dev_str=N
         ivy.permute_dims(
             ivy.astype(
                 ivy.linspace(
-                    ivy.zeros(new_img_dims, device=dev_str),
-                    ivy.ones(new_img_dims, device=dev_str) * (flat_batch_size - 1),
+                    ivy.zeros(new_img_dims, device=device),
+                    ivy.ones(new_img_dims, device=device) * (flat_batch_size - 1),
                     flat_batch_size,
                     axis=-1,
                 ),

diff --git a/ivy_vision/mesh.py b/ivy_vision/mesh.py
@@ -12,7 +12,7 @@
 
 
 def rasterize_triangles(
-    pixel_coords_triangles, image_dims, batch_shape=None, dev_str=None
+    pixel_coords_triangles, image_dims, batch_shape=None, device=None
 ):
     """Rasterize image-projected triangles based on:
     https://www.scratchapixel.com/lessons/3d-basic-rendering/rasterization-practical
@@ -29,7 +29,7 @@ def rasterize_triangles(
         Image dimensions.
     batch_shape
         Shape of batch. Inferred from Inputs if None. (Default value = None)
-    dev_str
+    device
         device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc.
         Same as x if None. (Default value = None)
 
@@ -42,8 +42,8 @@ def rasterize_triangles(
     if batch_shape is None:
         batch_shape = []
 
-    if dev_str is None:
-        dev_str = ivy.dev(pixel_coords_triangles)
+    if device is None:
+        device = ivy.dev(pixel_coords_triangles)
 
     # shapes as list
     batch_shape = list(batch_shape)
@@ -113,7 +113,7 @@ def rasterize_triangles(
         num_batch_dims_after = len(batch_dims_after)
 
         # [batch_dim]
-        batch_indices = ivy.arange(batch_dim, dtype="int32", device=dev_str)
+        batch_indices = ivy.arange(batch_dim, dtype="int32", device=device)
 
         # [1]*num_batch_dims_before x batch_dim x [1]*num_batch_dims_after x 1 x 1
         reshaped_batch_indices = ivy.reshape(
@@ -157,7 +157,7 @@ def rasterize_triangles(
     )
 
 
-def create_trimesh_indices_for_image(batch_shape, image_dims, dev_str=None):
+def create_trimesh_indices_for_image(batch_shape, image_dims, device=None):
     """Create triangle mesh for image with given image dimensions
 
     Parameters
@@ -166,7 +166,7 @@ def create_trimesh_indices_for_image(batch_shape, image_dims, dev_str=None):
         Shape of batch.
     image_dims
         Image dimensions.
-    dev_str
+    device
         device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc.
         (Default value = None)
 
@@ -188,20 +188,20 @@ def create_trimesh_indices_for_image(batch_shape, image_dims, dev_str=None):
 
     # 1 x W-1
     t00_ = ivy.reshape(
-        ivy.arange(image_dims[1] - 1, dtype="float32", device=dev_str), (1, -1)
+        ivy.arange(image_dims[1] - 1, dtype="float32", device=device), (1, -1)
     )
 
     # H-1 x 1
     k_ = (
         ivy.reshape(
-            ivy.arange(image_dims[0] - 1, dtype="float32", device=dev_str), (-1, 1)
+            ivy.arange(image_dims[0] - 1, dtype="float32", device=device), (-1, 1)
         )
         * image_dims[1]
     )
 
     # H-1 x W-1
-    t00_ = ivy.matmul(ivy.ones((image_dims[0] - 1, 1), device=dev_str), t00_)
-    k_ = ivy.matmul(k_, ivy.ones((1, image_dims[1] - 1), device=dev_str))
+    t00_ = ivy.matmul(ivy.ones((image_dims[0] - 1, 1), device=device), t00_)
+    k_ = ivy.matmul(k_, ivy.ones((1, image_dims[1] - 1), device=device))
 
     # (H-1xW-1) x 1
     t00 = ivy.expand_dims(t00_ + k_, axis=-1)
@@ -220,7 +220,7 @@ def create_trimesh_indices_for_image(batch_shape, image_dims, dev_str=None):
 
 
 def coord_image_to_trimesh(
-    coord_img, validity_mask=None, batch_shape=None, image_dims=None, dev_str=None
+    coord_img, validity_mask=None, batch_shape=None, image_dims=None, device=None
 ):
     """Create trimesh, with vertices and triangle indices, from co-ordinate image.
 
@@ -235,7 +235,7 @@ def coord_image_to_trimesh(
         Shape of batch. Inferred from inputs if None. (Default value = None)
     image_dims
         Image dimensions. Inferred from inputs in None. (Default value = None)
-    dev_str
+    device
         device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc.
         Same as x if None. (Default value = None)
 
@@ -245,8 +245,8 @@ def coord_image_to_trimesh(
         Vertices *[batch_shape,(hxw),3]* amd Trimesh indices *[batch_shape,n,3]*
 
     """
-    if dev_str is None:
-        dev_str = ivy.dev(coord_img)
+    if device is None:
+        device = ivy.dev(coord_img)
 
     if batch_shape is None:
         batch_shape = ivy.shape(coord_img)[:-3]
@@ -294,7 +294,7 @@ def coord_image_to_trimesh(
 
         # BS x 2x(H-1xW-1) x 3
         all_trimesh_indices = create_trimesh_indices_for_image(
-            batch_shape, image_dims, dev_str
+            batch_shape, image_dims, device
         )
 
         # BS x N x 3