Skip to content

Commit

Permalink
[Frontend] Support static_argnums (#476)
Browse files Browse the repository at this point in the history
**Context:** Supporting ``static_argnums`` to identify which arguments
are static. I

**Description of the Change:** Introduce the ``static_argnums``
parameter to ``@qjit`` and add a mechanism to determine if the function
needs to be recompiled when the static arguments are changed.

**Benefits:** Users can pass custom objects to a function with
``@qjit``. The ``QJIT`` object also stores all previously compiled
functions for different arguments. If the arguments are seen before,
there is no need for re-compilation.

**Possible Drawbacks:** The ``QJIT`` object will store all the
previously compiled functions and require more space.
The mechanism for recompilation is based on checking hash values, so
there might be a possibility of collision. It also creates a new
``catalyst.utils.filesystem.WorkspaceManager`` for every new function.
If we do not create a new ``WorkspaceManager``, the new function will
not be properly compiled. I am not sure if I should modify anything
related to it.

**Related GitHub Issues:** closes #461

---------

Co-authored-by: Tzung-Han Juang <tzung-han.juang@mail.mcgill.ca>
Co-authored-by: David Ittah <dime10@users.noreply.github.com>
  • Loading branch information
3 people committed Jan 31, 2024
1 parent b336b99 commit 96a8d1e
Show file tree
Hide file tree
Showing 6 changed files with 371 additions and 14 deletions.
39 changes: 39 additions & 0 deletions doc/changelog.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,44 @@

<h3>New features</h3>

* Catalyst now supports just-in-time compilation of static arguments.
[(#476)](https://github.com/PennyLaneAI/catalyst/pull/476)

The ``@qjit`` decorator can now be used to compile functions with static arguments with
the ``static_argnums`` keyword argument. ``static_argnums`` can be an integer or an iterable
of integers that specify which positional arguments of a compiled function should be treated as
static. This feature allows users to pass hashable Python objects to a compiled function (as the
static arguments).

A ``qjit`` object stores the hash value of a compiled function's static arguments. If any static
arguments are changed, the ``qjit`` object will check its stored hash values. If no hash value is
found, the function will be re-compiled with new static arguments. Otherwise, no re-compilation
will be triggered and the previously compiled function will be used.

```py
from dataclasses import dataclass

from catalyst import qjit

@dataclass
class MyClass:
val: int

def __hash__(self):
return hash(str(self))

@qjit(static_argnums=(1,))
def f(
x: int,
y: MyClass,
):
return x + y.val

f(1, MyClass(5))
f(1, MyClass(6)) # re-compilation
f(2, MyClass(5)) # no re-compilation
```

<h3>Improvements</h3>

* Add native support for `qml.PSWAP` and `qml.ISWAP` gates on Amazon Braket devices. Specifically, a circuit like
Expand Down Expand Up @@ -140,6 +178,7 @@ This release contains contributions from (in alphabetical order):
Mikhail Andrenkov,
Ali Asadi,
David Ittah,
Tzung-Han Juang,
Erick Ochoa Lopez,
Haochen Paul Wang.

Expand Down
Loading

0 comments on commit 96a8d1e

Please sign in to comment.