Only run required rewrites in JAX tests

Several tests ended up not testing the backend Op implementations due to constant folding of inputs.

Bumps [pypa/gh-action-pypi-publish](https://github.com/pypa/gh-action-pypi-publish) from 1.9.0 to 1.10.3.
- [Release notes](https://github.com/pypa/gh-action-pypi-publish/releases)
- [Commits](https://github.com/pypa/gh-action-pypi-publish/compare/v1.9.0...v1.10.3)

---
updated-dependencies:
- dependency-name: pypa/gh-action-pypi-publish
  dependency-type: direct:production
  update-type: version-update:semver-minor
...
Signed-off-by: dependabot[bot] <support@github.com>

* fixed merge conflicts

* fixed failing tests and added rewrite for pt.diag

* minor changes; added test to not apply rewrite

* added test for batched case and more cases of not applying rewrite

* minor changes

* Added rewrite for diag of kronecker product

* Added rewrite for slogdet; added docstrings for  rewrites

* fixed typo

* added rewrite for diag(block_diag)

* added rewrite for determinant of blockdiag

* Added rewrite for slogdet; added docstrings for all 3 rewrites

* fixed typecasting for tests

It was failing to detect the version

Co-authored-by: Ian Schweer <ischweer@riotgames.com>

https://github.com/pytorch/pytorch/issues/123097

Due to https://github.com/mamba-org/setup-micromamba/issues/225