Skip to content

Fix doctests in transforms for graph_updates=False default - #5441

Merged
copybara-service[bot] merged 1 commit into
google:mainfrom
samanklesaria:fix-transforms-doctest
May 12, 2026
Merged

copybara-service[bot] merged 1 commit into
google:mainfrom
samanklesaria:fix-transforms-doctest

Conversation

@samanklesaria

@samanklesaria samanklesaria commented Apr 24, 2026 •

Copy link
Copy Markdown
Collaborator

This PR makes the following changes in docstrings:

  • grad: remove DiffState usage in docstring
  • custom_vjp: add graph_updates=True to examples using input_updates_g bwd API, as custom_vjp fails with graph_updates=False. #5452 shows this is buggy otherwise.
  • vmap: remove use of StateAxes in docstring
  • optimizer: remove DiffState usage in docstring

- grad: add graph_updates=True where DiffState argnums or State output is expected
- custom_vjp: add graph_updates=True to examples using input_updates_g bwd API
- vmap: add graph_updates=True where StateAxes is used in in_axes
- optimizer: add graph_updates=True for DiffState argnums in grad call
@samanklesaria
samanklesaria force-pushed the fix-transforms-doctest branch from eb45cb1 to bd6bf36 Compare April 30, 2026 16:27
@copybara-service
copybara-service Bot merged commit e9aaebe into google:main May 12, 2026
21 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants