I had a look at the broken apidocs. Unfortunately I missed that there already was a branch / PR about that. Feel free to decide if this looks good to you, or to ignore this.
On my system, the apidocs are not properly generated on the master branch, but with the changes here it works. I don't know how to generate the readthedocs to try whether it works there.
Thanks @pmenczel. I have made similar changes on the other branch. After fixing some additional errors (due to JAX updates), I will update the main branch.
I had a look at the broken apidocs. Unfortunately I missed that there already was a branch / PR about that. Feel free to decide if this looks good to you, or to ignore this.
On my system, the apidocs are not properly generated on the master branch, but with the changes here it works. I don't know how to generate the readthedocs to try whether it works there.