[TF] Refactor gradients computation with expandingShape (#24249)
* Refactor with new expandingShape function
* Add space
Co-Authored-By: sgugger <35901082+sgugger@users.noreply.github.com>
* Add space
Co-Authored-By: sgugger <35901082+sgugger@users.noreply.github.com>
