DOC: pm.sampling.jax functions are not included in PyMC docs

[cluhmann]

Issue with current documentation:

Various functions made available via pm.sampling.jax, including sample_numpyro_nuts() and sample_blackjax_nuts() do not appear in the documentation despite these functions having docstrings. Maybe we didn't include the pm.sampling.jax module is some auto-docing because it was a bit experimental at the time?

[OriolAbril]

They are here: https://www.pymc.io/projects/docs/en/stable/api/samplers.html#samplers. Should they have their own section instead of being listed along other "sampling" functions?

[cluhmann]

I guess the issue is that they don't show up in search results?

[OriolAbril]

They don't? What was the search string?

[cluhmann]

Here is an example. The closest thing that comes back is a link to the entry in the index.

Edit: numpyro too.

[OriolAbril]

it looks like the search engine works but is quite bad 😓. Not sure how to fix this. It looks like sample_numpyro_nuts is treated as a substring of the pymc.sampling.jax.sample_numpyro_nuts string, and substring search is a bit limited.

First of all, I am not sure if the search functionality is provided by sphinx, the theme or by readthedocs. In addition, as we have multiple builds and libraries, the search defaults to searching over all subprojects (which seems to add some extra challenges).

• Example of working substring search: https://www.pymc.io/search.html?q=umbe. Finds all occurrences of number but only in the homepage+blog website. https://www.pymc.io/projects/docs/en/stable/search.html?q=umbe shows many more results that do not appear in the previous one for some reason unknown to me.
• Example of not working substring search: https://www.pymc.io/search.html?q=sample_blackjax_nuts as you shared above. The first points to the source code for the function because it is not considered a substring there. Even if we search only in the learn and api sections, it doesn't find anything: https://www.pymc.io/projects/docs/en/stable/search.html?q=sample_blackjax_nuts. However, if we use what is considered the whole string: https://www.pymc.io/search.html?q=pymc.sampling.jax.sample_blackjax_nuts

It would be great if someone could look into how the search function works on the website. I'll be more than happy to answer any question about the website structure and building process