PennyLaneAI / PennyLaneAI/catalyst

[BUG] [FRONTEND] nested autograph functions results in autograph warnings

Open
#386 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug
Dominant language
Python
Stars
234
Forks
84
Avg merge
2d 15h
Merged PRs (30d)
66

Description

Shortcut: https://app.shortcut.com/xanaduai/story/51540/nested-autograph-conversions-fail

When using nested functions with autograph=True, various autograph warnings will be raised when calling the outermost function. However, autograph still appears to be working correctly.

For example, consider:

@qjit(autograph=True)
def f(x):
    if x > 5:
        y = x ** 2
    else:
        y = x ** 3
    return y

@qjit(autograph=True)
def g(x, n):
    for i in range(n):
        x = x + f(x)
    return x

Calling the inner function works fine:

>>> f(0.4)
array(0.064)

However, calling the outermost function will result in warnings:

>>> g(0.4, 6)
WARNING:tensorflow:AutoGraph could not transform <function outer_factory.<locals>.inner_factory.<locals>.f_1 at 0x7ad76c5b3ac0> and will run it as-is.
Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.
Cause: closure mismatch, requested ('ag__',), but source function had ()
To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert
WARNING: AutoGraph could not transform <function outer_factory.<locals>.inner_factory.<locals>.f_1 at 0x7ad76c5b3ac0> and will run it as-is.
Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.
Cause: closure mismatch, requested ('ag__',), but source function had ()
To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert
array(22.14135448)

Examining the jaxpr appears to indicate that Autograph is working correctly:

>>> g.jaxpr
{ lambda ; a:f64[] b:i64[]. let
    c:f64[] = for_loop[
      apply_reverse_transform=False
      body_jaxpr={ lambda ; d:i64[] e:f64[]. let
          f:bool[] = gt e 5.0
          g:f64[] = cond[
            branch_jaxprs=[{ lambda ; a:f64[] b:f64[]. let c:f64[] = integer_pow[y=2] a in (c,) }, { lambda ; a:f64[] b:f64[]. let c:f64[] = integer_pow[y=3] b in (c,) }]
          ] f e e
          h:f64[] = add e g
        in (h,) }
      body_nconsts=0
    ] 0 b 1 0 a
  in (c,) }

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reproducing the nested qjit(autograph=True) example from the issue and inspect the frontend path that handles nested Autograph conversions. Verify the warning is eliminated while the outer function still returns the shown result and produces a correct jaxpr.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, tensorflow
Domain
frontend
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.