Add jax implementation of pt.linalg.pinv

[jessegrabowski]

Motivation for these changes

Add a jax implementation of pt.linalg.pinv

Implementation details

Not the most elaborate PR:

@jax_funcify.register(MatrixPinv)
def jax_funcify_Pinv(op, **kwargs):
    def pinv(x):
        return jnp.linalg.pinv(x)

    return pinv

Checklist

• [ x] Explain motivation and implementation 👆
• [ x] Make sure that the pre-commit linting/style checks pass.
• Link relevant issues, preferably in nice commit messages.
• [x ] The commits correspond to relevant logical changes. Note that if they don't, we will rewrite/rebase/squash the git history before merging.
• [x ] Are the changes covered by tests and docstrings?
• Fill out the short summary sections 👇

Major / Breaking Changes

None

New features

You can compile graphs with pinv to JAX

Bugfixes

None

Documentation

None

Maintenance

None

[ricardoV94]

Thanks @jessegrabowski

[ricardoV94]

Oops I missed one float32 test that was failing: https://github.com/pymc-devs/pytensor/actions/runs/4968627533/jobs/8891362025?pr=294

[jessegrabowski]

I didn't add astype(config.floatX) to the test array (again...). Not sure how to proceed -- I deleted the branch on my fork already so it's not auto-pushing the change into this PR.

[ricardoV94]

No worries, I opened a fix PR in #296