trax というフレームワークを使う。 github ドキュメント

trax.math.grad : 数式処理してくれる