Add static_argnames option to qjit

doc/releases/changelog-dev.md

@@ -147,6 +147,30 @@
   Array([2, 4, 6], dtype=int64)
   ```
 
+* Static arguments of a jit-ted function can now be indicated by a `static_argnames`
+  argument to `qjit`.
+  [(#1158)](https://github.com/PennyLaneAI/catalyst/pull/1158)
+
+  ```python
+  @qjit(static_argnames="y")
+  def f(x, y):
+    if y < 10:  # y needs to be marked as static since its concrete boolean value is needed
+        return x + y
+
+  @qjit(static_argnames=["x","y"])
+  def g(x, y):
+    if x < 10 and y < 10:
+        return x + y
+
+  res_f = f(1, 2)
+  res_g = g(3, 4)
+  print(res_f, res_g)
+  ```
+
+  ```pycon
+  >>> 3 7
+  ```
+
 <h3>Improvements</h3>
 
 * Scalar tensors are eliminated from control flow operations in the program, and are replaced with

frontend/catalyst/jit.py

@@ -81,6 +81,7 @@ def qjit(
     logfile=None,
     pipelines=None,
     static_argnums=None,
+    static_argnames=None,
     abstracted_axes=None,
     disable_assertions=False,
     seed=None,
@@ -123,6 +124,8 @@ def qjit(
             considered to be used by advanced users for low-level debugging purposes.
         static_argnums(int or Seqence[Int]): an index or a sequence of indices that specifies the
             positions of static arguments.
+        static_argnames(str or Seqence[str]): a string or a sequence of strings that specifies the
+            names of static arguments.
         abstracted_axes (Sequence[Sequence[str]] or Dict[int, str] or Sequence[Dict[int, str]]):
             An experimental option to specify dynamic tensor shapes.
             This option affects the compilation of the annotated function.
@@ -430,6 +433,21 @@ def sum_abstracted(arr):
     if fn is None:
         return functools.partial(qjit, **kwargs)
 
+    # Map static_argnames to static_argnums
+    # no need to propagate static_argnames further
+    if static_argnames is not None:
+        static_argnums = [] if (static_argnums is None) else list(kwargs["static_argnums"])
+        fn_argnames = inspect.getfullargspec(fn).args
+        for argname in fn_argnames:
+            if argname in static_argnames:
+                static_argnums.append(fn_argnames.index(argname))
+
+        # Remove potential duplicates from static_argnums and static_argnames
+        static_argnums = list(dict.fromkeys(static_argnums))
+        static_argnums.sort()
+        kwargs["static_argnums"] = static_argnums
+    kwargs.pop("static_argnames")
+
     return QJIT(fn, CompileOptions(**kwargs))
 
 

frontend/test/pytest/test_static_arguments.py

@@ -18,6 +18,7 @@
 
 import pennylane as qml
 import pytest
+from jax.errors import TracerBoolConversionError
 
 from catalyst import qjit
 from catalyst.utils.exceptions import CompileError
@@ -221,6 +222,55 @@ def wrapper(x, c):
         captured = capsys.readouterr()
         assert captured.out.strip() == "Inside QNode: 0.5"
 
+    def test_static_argnames(self):
+        # pylint: disable=inconsistent-return-statements
+        """Test static arguments specified by names"""
+
+        @qjit(static_argnames="y")
+        def f(x, y):
+            if y < 10:
+                return x + y
+
+        assert f(1, 2) == 3
+
+        @qjit(static_argnames="x")
+        def g(x, y):
+            if y < 10:
+                return x + y
+
+        with pytest.raises(
+            TracerBoolConversionError, match="Attempted boolean conversion of traced"
+        ):
+            g(1, 2)
+
+        @qjit(static_argnames=("x", "y"))
+        def h(x, y):
+            if x < 10 and y < 10:
+                return x + y
+
+        assert h(1, 2) == 3
+
+        @qjit(static_argnames=("x"), static_argnums=[1])
+        def p(x, y):
+            if x < 10 and y < 10:
+                return x + y
+
+        assert p(1, 2) == 3
+
+        @qjit(static_argnames=("y"), static_argnums=[0])
+        def q(x, y):
+            if x < 10 and y < 10:
+                return x + y
+
+        assert q(1, 2) == 3
+
+        @qjit(static_argnames=("y"), static_argnums=[1])
+        def r(x, y):
+            if y < 10:
+                return x + y
+
+        assert r(1, 2) == 3
+
 
 if __name__ == "__main__":
     pytest.main(["-x", __file__])