Fix nested accelerate with grad

[erick-xanadu]

Context: Fixes taking the gradient of nested accelerate calls.

Description of the Change: Takes advantage of accelerate context to let JAX do its thing.

Benefits: Fixes bug.

Possible Drawbacks: I think a better way to do this would be to have is_tracing to return False. However, that would involve refactoring the current tracing context. Maybe a possible follow up.

Related GitHub Issues: Fixes #1086

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 97.86%. Comparing base (fe5e6ae) to head (0130660).
Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1156   +/-   ##
=======================================
  Coverage   97.86%   97.86%           
=======================================
  Files          76       76           
  Lines       10811    10811           
  Branches     1281     1281           
=======================================
  Hits        10580    10580           
  Misses        179      179           
  Partials       52       52           

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

[dime10]

Thanks Erick :)