Added trace alias to jnp.linalg

[vfdev-5]

Related to #21088

cc @Micky774

[Micky774]

Instead of adding an alias in the linalg source file, could we just add a direct import of the source of jnp.trace into the linalg namespace file?

[vfdev-5]

Instead of adding an alias in the linalg source file, could we just add a direct import of the source of jnp.trace into the linalg namespace file?

I have no strong opinion on the approach. Yes, we can do that as well. In the current approach we sort of hidden that detail into to the _src.

[jakevdp]

Sorry, just realized that this new function is missing from the docs here: https://github.com/google/jax/blob/main/docs/jax.scipy.rst#jaxscipylinalg

[vfdev-5]

Sorry, just realized that this new function is missing from the docs here: https://github.com/google/jax/blob/main/docs/jax.scipy.rst#jaxscipylinalg

I should be missing something as I'm not sure to understand your point here. As far as I understand, there is no jax.scipy.linalg.trace function and so no docs for it.

[jakevdp]

Sorry, I meant to point to the numpy docs: https://github.com/google/jax/blob/main/docs/jax.numpy.rst#jaxnumpylinalg

We need to add the function to that list, or it won't show up in the API docs on the website.