Add jax.scipy.linalg.toeplitz.

[yotarok]

A counterpart for "scipy.linalg.toeplitz" is added in "jax.scipy.linalg".

We recently have a case that requires a Jax counterpart of tf.linalg.LinearOperatorToeplitz, and this functionality is in fact a part of SciPy, implemented as scipy.linalg.toeplitz.
This also relates to the request in #10144.

[jakevdp]

Looks great! Just a few comments below and we can get this in.

Also, please add the new function to the list here so it will show up in the HTML docs: https://github.com/google/jax/blob/main/docs/jax.scipy.rst#jaxscipylinalg

[yotarok]

Thanks for reviewing.
All the issues are resolved and the new version is force-pushed to the branch.

[jakevdp]

Looks great, thanks!

[yotarok]

Looks like the numerical error is severe when TPUs/ GPUs are used because lax.conv_general_dilated_patches is used.
So, I switched to use vmap and np.roll that are slightly slower but should be accurate.

[hawkinsp]

@yotarok You can get higher precision on TPU/GPU with jax.default_matmul_precision("float32") which can be used as, e.g., a decorator, or the precision option to the convolution.

[jakevdp]

I suspect vmap of roll is going to be much less efficient than the original approach, due to how it's implemented. Would you like to change back to the convolution-based approach along with adjusting the matmul precision?

[yotarok]

The profiling results with TPU were:

• lax.conv_general_dilated_patches: ~9ms
• roll + vmap: ~15ms
• lax.conv_general_dilated_patches + Precision.HIGHEST: ~12ms
  (with configuaration of r == c, r.shape = (1024, 256) and toeplitz is called with vmap).

Confirmed that with Precision.HIGHEST, assertAllClose in the test passes with the default tolerance parameters with TPU/ GPU.

[yotarok]

I noticed that integers are not tested. So, I added integer tests and some skip conditions because CuDNN doesn't support integer convolution except for int8.
PTAL again.

[jakevdp]

Last couple minor comments - thanks!

[jakevdp]

Here and throughout, please use two-space indentations.,

[yotarok]

done.
What's the style recommendation for hanging indents?
For _src/scipy/linalg.py, I followed Google-style, i.e. 2 spaces for normal indents, and 4 spaces for hanging indents.
However, tests/linalg_test.py uses 2 spaces also for hanging indents, so I followed that file-local convention only for the test.

[yotarok]

Thank you for the comments, I updated the PR accordingly.

[jakevdp]

Thanks!

[yotarok]

Ping, since it's been while after the PR was approved, the branch is rebased on top of the current main.