[Frontend] Add support for while loops in autograph

[sergei-mironov]

Add support for while-loops in the @qjit(autograph=True) mode

We replace the while-loops with catalyst.while_loop calls on a best-efforts basis. In case of errors we fallback to the unrolled loop path. A decision was made to not warn by default that the fallback has been used.

[sc-41300]

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (9971a05) 99.56% compared to head (35646fa) 99.62%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #318      +/-   ##
==========================================
+ Coverage   99.56%   99.62%   +0.05%     
==========================================
  Files          42       42              
  Lines        7410     7438      +28     
  Branches      433      434       +1     
==========================================
+ Hits         7378     7410      +32     
+ Misses         16       14       -2     
+ Partials       16       14       -2     

Files	Coverage Δ
frontend/catalyst/ag_primitives.py	100.00% <100.00%> (+2.07%)	⬆️

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

[dime10]

Thanks @grwlf, looking good so far 🚀
Some preliminary comments:

[dime10]

Oh, this PR will also need to update the newly added autograph docs :)

[josh146]

Oh, this PR will also need to update the newly added autograph docs :)

I agree, but in this case might argue for doing it as a separate PR - just because the docs still pin to latest, so readers will be seeing support for something that they can't pip install.

[dime10]

I agree, but in this case might argue for doing it as a separate PR - just because the docs still pin to latest, so readers will be seeing support for something that they can't pip install.

This is tricky, because we can't prevent new functions from not showing up in the docs either 🤔 (well we probably could but ...)
I think that's why I originally wanted the docs to show the latest release.
@josh146 Do you think it would make sense to have a branch that's identical to the last release, but can contain corrections (or additions like the autograph docs), and have the docs show that branch by default?
Then latest could be kept in sync with the current code status in main.

[josh146]

@josh146 Do you think it would make sense to have a branch that's identical to the last release, but can contain corrections (or additions like the autograph docs), and have the docs show that branch by default?
Then latest could be kept in sync with the current code status in main.

My recommendation would be:

• We begin pinning our docs to stable
• We do 'post' releases (e.g., 0.3.1-post1, 0.3.1.post2, etc.) on GitHub whenever we need to update stable docs (this usually only happens very rarely).
  • Post releases do not include any logical changes, just doc changes.
  • These releases are only made on GitHub, not PyPI.

On read the docs, stable will automatically follow the most recent github release, whether post or not.

[dime10]

@grwlf Can you open a PR that updates the autograph docs (to be merged later)?

[sergei-mironov]

@grwlf Can you open a PR that updates the autograph docs (to be merged later)?

@dime10 here is the doc PR #330

[dime10]

Thanks @grwlf, just a few points to resolve and we can merge this 👍

[dime10]

Great work 💯