Jax implementation of log1mexp

[kylejcaron]

This PR creates a Jjax implementation of the log1mexp op, to resolve issue #990

Thank you for opening a PR!

Here are a few important guidelines and requirements to check before your PR can be merged:

• There is an informative high-level description of the changes.
• The description and/or commit message(s) references the relevant GitHub issue(s).
• pre-commit is installed and set up.
• The commit messages follow these guidelines.
• The commits correspond to relevant logical changes, and there are no commits that fix changes introduced by other commits in the same branch/BR.
• There are tests covering the changes introduced in the PR.

Don't worry, your PR doesn't need to be in perfect order to submit it. As development progresses and/or reviewers request changes, you can always rewrite the history of your feature/PR branches.

If your PR is an ongoing effort and you would like to involve us in the process, simply make it a draft PR.

[kylejcaron]

quick callout: pre-commit had the following failure unrelated to any code introduced in my PR so I uninstalled pre-commit and proceeded to commit my work

mypy.....................................................................Failed
- hook id: mypy
- exit code: 1

aesara/scan/op.py:3431: error: Unused "type: ignore" comment
Found 1 error in 1 file (checked 2 source files)

[ricardoV94]

Looks good, I have a suggestion below for a more thorough check

[ricardoV94]

quick callout: pre-commit had the following failure unrelated to any code introduced in my PR so I uninstalled pre-commit and proceeded to commit my work

mypy.....................................................................Failed
- hook id: mypy
- exit code: 1

aesara/scan/op.py:3431: error: Unused "type: ignore" comment
Found 1 error in 1 file (checked 2 source files)

I have seen that can sometimes happen when the pre-commit runs only on the modified files but goes away when it runs on all files. It's annoying though

[codecov]

Codecov Report

Merging #994 (f3c430a) into main (2ccd9cc) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##             main     #994   +/-   ##
=======================================
  Coverage   79.23%   79.23%           
=======================================
  Files         152      152           
  Lines       47946    47951    +5     
  Branches    10916    10916           
=======================================
+ Hits        37990    37995    +5     
  Misses       7448     7448           
  Partials     2508     2508           

Impacted Files	Coverage Δ
aesara/link/jax/dispatch.py	81.89% <100.00%> (+0.15%)	⬆️

[ricardoV94]

Looks good, can you squash the three commits into a single one?

[kylejcaron]

Looks good, can you squash the three commits into a single one?

just squashed into 1 commit!

[ricardoV94]

Thanks @kylejcaron!

[kylejcaron]

Thanks @kylejcaron!

Thanks for the help on this!