BUG: missing `track_gradient` in some functions

[Xylambda]

• https://github.com/Xylambda/toydiff/blob/master/src/toydiff/core.py#L506
• https://github.com/Xylambda/toydiff/blob/master/src/toydiff/core.py#L362