Fix doctests in transforms for graph_updates=False default
- 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
