linting

PennyLaneAI · Sep 23, 2024 · a40eb46 · a40eb46
1 parent 5754b87
commit a40eb46
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/frontend/catalyst/api_extensions/function_maps.py b/frontend/catalyst/api_extensions/function_maps.py
@@ -320,10 +320,11 @@ def _get_batch_size(self, args_flat, axes_flat, axis_size):
 
         batch_sizes = []
         for i, (arg, d) in enumerate(zip(args_flat, axes_flat)):
-            shape = np.shape(arg)
             if d is None:
                 continue
-            elif len(shape) > d:
+
+            shape = np.shape(arg)
+            if len(shape) > d:
                 batch_sizes.append(shape[d])
             else:
                 raise ValueError(

diff --git a/frontend/test/pytest/test_vmap.py b/frontend/test/pytest/test_vmap.py
@@ -21,6 +21,8 @@
 
 from catalyst import qjit, vmap
 
+# pylint: disable=bug-vmap-shape
+
 
 class TestVectorizeMap:
     """Test QJIT compatibility with JAX vectorization."""