Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 97.95%. Comparing base (b12e047) to head (d9774f5).
Report is 1 commits behind head on main.

@@           Coverage Diff           @@
##             main    #1295   +/-   ##
=======================================
  Coverage   97.95%   97.95%           
=======================================
  Files          77       77           
  Lines       11318    11320    +2     
  Branches      981      981           
=======================================
+ Hits        11087    11089    +2     
  Misses        181      181           
  Partials       50       50           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

Hi @erick-xanadu, thank you for these comments. Both the points about using pytest.mark.parametrize and LLVM's lit functionality are very insightful, and I will keep these in mind for future development.

On your other comment about calling the function using keyword arguments, I did some debugging and research. It looks like this is an issue independent of having default arguments. I made some simple test cases which do not use default arguments, and these fail on using keyword arguments too. It appears as though the static_argnums=... option in the decorator does not apply if the index refers to a keyword argument resulting in a error like jax.errors.TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].. (if the argument in question is used in a boolean expression).

It appears as though the static_argnums=... option in the decorator does not apply if the index refers to a keyword argument

Interesting, can you provide a small example in the comments? Does it only fail with qjit, or does jax.jit fail as well?

Interesting, can you provide a small example in the comments? Does it only fail with qjit, or does jax.jit fail as well?

Yes, and can you provide a bit more explanation as to why it fails?

Hi @paul0403 and @erick-xanadu,

So here's an example of a case that fails

        @qjit(static_argnums=[1])
        def f(y, x):
            if x < 10:
                return x + y
            return 42000

        assert f(20, x=3) == 23

This test case, however, passes when we use jax.jit, like in the example below

        @partial(jax.jit, static_argnums=[1])
        def f(y, x):
            if x < 10:
                return x + y
            return 42000

        assert f(20, x=3) == 23

Upon doing some step through debugging, I realized that this happens because capture() only filters static arguments from args and not from kwargs, and in the above case the the static argument is in kwargs. I have added comments in the function snippet below to pinpoint the issue.

    def capture(self, args, **kwargs):
        """Capture the JAX program representation (JAXPR) of the wrapped function.

        Args:
            args (Iterable): arguments to use for program capture

        Returns:
            ClosedJaxpr: captured JAXPR
            PyTreeDef: PyTree metadata of the function output
            Tuple[Any]: the dynamic argument signature
        """

        # In the above case `args=(20)` and `kwargs={'x': 3}`.
        # The below lines of code only process static arguments from `args`.

        verify_static_argnums(args, self.original_function, self.compile_options.static_argnums)
        static_argnums = self.compile_options.static_argnums
        abstracted_axes = self.compile_options.abstracted_axes

        dynamic_args = filter_static_args(args, static_argnums)
        dynamic_sig = get_abstract_signature(dynamic_args)
        full_sig = merge_static_args(dynamic_sig, args, static_argnums)
        ...

@AniketDalvi , thanks for your answer! Will you be implementing a solution to the test case regarding keyword arguments or are you done with your submission? Cheers!

Hi! I am curious to understand how this would affect my assessment. It would be a slightly larger undertaking as it appears that the keyword argument test case doesn't just pertain to #1163. Although I would be glad to take a stab at it, I am currently also tied up in some PhD deadlines.

Furthermore, I recently also notified the recruiter that I received a verbal offer from a company I was interviewing with and requested to expedite the recruitment process as this role with Xanadu is a top priority for me.

Hi @AniketDalvi, don't worry about it. Submission looks good. We will contact you back by e-mail. Cheers. :)