no cover loop dynamic shape part

Author: JerryChen97

pennylane/control_flow/_loop_abstract_axes.py

@@ -29,7 +29,7 @@
 AbstractShapeLocation = namedtuple("AbstractShapeLocation", ("arg_idx", "shape_idx"))
 
 
-def add_abstract_shapes(f, shape_locations: list[list[AbstractShapeLocation]]):
+def add_abstract_shapes(f, shape_locations: list[list[AbstractShapeLocation]]):  # pragma: no cover
     """Add the abstract shapes at the specified locations to the output of f.
 
     Here we can see that the shapes at argument 0, shape index 0 and
@@ -65,7 +65,7 @@ def new_f(*args, **kwargs):
     return new_f
 
 
-def get_dummy_arg(arg):
+def get_dummy_arg(arg):  # pragma: no cover
     """If any axes are abstract, replace them with an empty numpy array.
 
     Even if abstracted_axes specifies two dimensions as having different dynamic shapes,
@@ -112,7 +112,7 @@ def validate_no_resizing_returns(
     """
     offset = len(locations)  # number of abstract shapes. We start from the first normal arg.
 
-    for locations_list in locations:
+    for locations_list in locations:  # pragma: no cover
         loc0 = locations_list[0]
         first_var = jaxpr.outvars[loc0.arg_idx + offset].aval.shape[loc0.shape_idx]
         for compare_loc in locations_list[1:]:
@@ -132,7 +132,7 @@ def validate_no_resizing_returns(
 
 
 def _has_dynamic_shape(val):
-    return any(not isinstance(s, int) for s in getattr(val, "shape", ()))
+    return any(not isinstance(s, int) for s in getattr(val, "shape", ()))  # pragma: no cover
 
 
 def handle_jaxpr_error(
@@ -142,7 +142,7 @@ def handle_jaxpr_error(
     about 'Incompatible shapes for broadcasting'."""
     import jax  # pylint: disable=import-outside-toplevel
 
-    if "Incompatible shapes for broadcasting" in str(e) and jax.config.jax_dynamic_shapes:
+    if "Incompatible shapes for broadcasting" in str(e) and jax.config.jax_dynamic_shapes:  # pragma: no cover
         closures = sum(((fn.__closure__ or ()) for fn in fns), ())
         if any(_has_dynamic_shape(i.cell_contents) for i in closures):
             msg = (
@@ -176,7 +176,7 @@ def add_arg(self, x_idx: int, x):
         arg_abstracted_axes = {}
 
         for shape_idx, s in enumerate(getattr(x, "shape", ())):
-            if not isinstance(s, int):  #  if not int, then abstract
+            if not isinstance(s, int):  #  pragma: no cover
                 found = False
                 if not self.allow_array_resizing:
                     for previous_idx, previous_shape in enumerate(self.abstract_shapes):
@@ -189,7 +189,7 @@ def add_arg(self, x_idx: int, x):
                             break
                 # haven't encountered it, so add it to abstract_axes
                 # and use new number designation
-                if not found:
+                if not found:  # pragma: no cover
                     arg_abstracted_axes[shape_idx] = len(self.abstract_shapes)
                     self.shape_locations.append([AbstractShapeLocation(x_idx, shape_idx)])
                     self.abstract_shapes.append(s)
@@ -262,5 +262,5 @@ def f(*args, allow_array_resizing):
     if not any(calculator.abstracted_axes):
         return None, [], []
 
-    abstracted_axes = jax.tree_util.tree_unflatten(structure, calculator.abstracted_axes)
-    return abstracted_axes, calculator.abstract_shapes, calculator.shape_locations
+    abstracted_axes = jax.tree_util.tree_unflatten(structure, calculator.abstracted_axes)  # pragma: no cover
+    return abstracted_axes, calculator.abstract_shapes, calculator.shape_locations  # pragma: no cover