ENH: Expansion of pymc.math to include more numpy/pytensor functions

[brandonhorsley]

Before

Due to pymc.math missing a number of common numpy/pytensor operations, this is primarily the inverse trig functions but it would be worth combing through the full list to see if there are any other notable absences. However since these functions can still be found through pytensor like so:


import pymc as pm
import pytensor.tensor as pt

test=pt.scalar(name="test")
res=pm.math.arccos(test) #Fails
res=pt.arccos(test) #Succeeds

### After

```python
This should work essentially the same as the other pymc.math functions so along the lines of:


import pymc as pm
import pytensor.tensor as pt

test=pt.scalar(name="test")
res=pm.math.arccos(test) #Succeeds

### Context for the issue:

This issue was spawned from this post from the PyMC discourse (https://discourse.pymc.io/t/no-inverse-trig-in-pm-math-functions/13681/2) and the slightly later discussion post (https://github.com/pymc-devs/pymc/discussions/7124). The issue is mainly a matter of convenience so as to have a complete set of numpy/pytensor operations all accessible through pm.math, and the pymc.math documentation on the website should hopefully reflect this as well.

[welcome]

]
🎉 Welcome to PyMC! 🎉 We're really excited to have your input into the project! 💖

If you haven't done so already, please make sure you check out our Contributing Guidelines and Code of Conduct.

[ricardoV94]

@brandonhorsley do you want to add any obvious missing functions? Feel free to open a PR for it

[brandonhorsley]

Basic Answer

Apologies for taking a while to get round to this, I am happy enough to have a crack at it and I have done some figuring out to get a basic starting list of functions that need to be sorted, secondly a list of functions that work but aren't mentioned and after I will mention some additional functions that may be worth an inclusion.

Functions that are mentioned in all but aren't mentioned in documentation and don't seem to work when called:

• arccos
• arccosh
• arcsin
• arcsinh
• arctan
• arctanh
• broadcast_to
• cumprod
• cumsum
• full_like
• matmul
• max
• mean
• min
• round

The only one that isn't a pytensor.tensor import in math.py (file linked in boring derivation below) is 'round' and is already defined in the file. It not working for me may just be an issue with me specifically, tround does work which is supposed to be deprecated so I may have an out-of-date version, need to investigate.

Functions that are defined in math.py that do work but aren't included in website documentation so can be added now if needed:

• erfc
• erfcinv
• log1pexp
• logsumexp
• kronecker
• cartesian
• kron_dot
• kron_solve_lower
• kron_solve_upper
• kron_diag
• flat_outer
• logdiffexp
• logdiffexp_numpy
• softmax
• log_softmax
• logbern
• log1mexp
• log1mexp_numpy
• flatten_list
• logdet
• batched_diag
• block_diagonal
• tround (noted as deprecated though so should likely be omitted until round is sorted)

Personal opinion on maths functions that could be worthy additions would be:

• argmax (already found on pytensor.tensor)
• einsum?
• adjoint
• transpose

Boring derivation reasoning

Cross-referenced the list on all in https://github.com/pymc-devs/pymc/blob/main/pymc/math.py#L100 to ones that seem to be missing from the pymc documentation page (https://www.pymc.io/projects/docs/en/stable/api/math.html). Then after noticing that some in this list were defined in the math.py file I checked through the list of all and any bonuses that were mentioned in math.py,eliminated [ones,zeros,full] since ones_like and zeros_like have already been chosen, then I called each in the remaining list in a python window, any that raised the following error were noted as needing to be fixed whilst those that didn't were noted as needing to be added to the website documentation:
AttributeError: module 'pymc.math' has no attribute 'VarName'

Finally consulting pytensor.tensor page, numpy page and pytorch was used to produce a small list of possibly worthy ideas for bonus functions to add.

Questions/points of note

As a bit of a noob concerning more of this stuff it seems curious that only a select few of the pytensor.tensor imports seem not to work since they are also included in all. So abs is imported from pytensor and included in all and works as intended, but arccos and the like are also imported from pytensor and included in all but don't work. My understanding is that surely all the entries in all should be imported into as pymc.math or is there a different file I should be looking at?

Conscientious of not suggesting too many bonus functions to add to keep pymc running speedy.

Unsure of conduct where this documentation issue has popped up as part of looking into this, as to whether someone will add a documentation tag to this or to raise a separate issue with the documentation tag.

[brandonhorsley]

Update, in running a clean install in order to work on a pull request, the functions that were bugged for me work perfectly on the clean environment so the source of that must have been unique to my old anaconda environment. Ergo this simplifies the rest of this issue down to adding all of these functions from the earlier post to be added to the documentation which shouldn't take long at all. Additional functions that weren't in the perceived to be bugged list can still be added too in the pull request that I am working on.

[ricardoV94]

@brandonhorsley are you still interested in adding the missing entries to the docs?

[brandonhorsley]

Hi @ricardoV94 yes I am, I just got stuck on trying to setup a way to preview the documentation on my Windows machine, I made edits to the math.rst file for a pull request which I think should hopefully work but without a way to actually preview the changes I didn't want to do the pull request until I could check that it would actually work! If possible then I could do with some help setting it up in order to do that check and then do the pull request

[ricardoV94]

Hi @ricardoV94 yes I am, I just got stuck on trying to setup a way to preview the documentation on my Windows machine, I made edits to the math.rst file for a pull request which I think should hopefully work but without a way to actually preview the changes I didn't want to do the pull request until I could check that it would actually work! If possible then I could do with some help setting it up in order to do that check and then do the pull request

If you open a PR there is a remote preview of the docs. It's a bit slower to iterate, but if it worked from the first try you would be all set :)

[brandonhorsley]

As in doing the github preview? as in,

So i don't think it would lead to any issues since as far as I can tell due to the way autosummary functionality but I fear I am missing something in seeing how it would look on the website equivalent with the divisions and the description pulled from the relevant module:

[ricardoV94]

Nope, when you open a PR, a link will be automatically added at the end of the first post to the preview of the docs (takes a while to show up):

For example: #7206

https://pymcio--7206.org.readthedocs.build/projects/docs/en/7206
https://pymcio--7206.org.readthedocs.build/projects/docs/en/7206/api.html

[brandonhorsley]

Ahh OK, in which case i will do the pull request and cross my fingers I got it right first try, thanks!

[ricardoV94]

Ahh OK, in which case i will do the pull request and cross my fingers I got it right first try, thanks!

It's fine if it's not right first try. If you push changes it will re-render again. It's just a very slow way to iterate, but hopefully you won't need too much

[brandonhorsley]

Still it at least gives me a way to iterate, I have started a draft pull request and I can see what you were referring to and am waiting for the build now! Much appreciated!

[DavidRavnsborg]

I would very much like to see einsum added to pytensor functions. I am currently stuck trying to fit a matrix of coefficients that are multiplied by a 3D tensor, and it would be ideal to be able to go:

contribution = pt.einsum("ijk,kl->ij", tensor_product, coef_matrix)

I currently do this with numpy, but I can't fit coefficients with numpy:

contribution = np.einsum("ijk,kl->ij", tensor_product, coef_matrix)

[ricardoV94]

Einsum is already present in pytensor, you may need to bump pymc.