Jake Vanderplas

Hi - this is more-or-less working as expected, because pytree equivalence includes equivalence of static elements. Can you say more about why it's not possible to specify `in_axes` with an...

I don't entirely follow – if `None` can be used in place of full subtrees, how could you possibly expect a generic `in_axis` to match a runtime pytree that may...

Hi @joamatab – we're chatting offline about how to best resolve this. In the meantime, I'm curious to hear about your experience here: the relevant APIs here have been raising...

I see a deprecation warning in JAX 0.4.23: ``` >>> import jax >>> print(jax.__version__) 0.4.23 >>> jax.interpreters.xla.backend_specific_translations :1: DeprecationWarning: jax.interpreters.xla.backend_specific_translations is deprecated. Register custom primitives via jax.interpreters.mlir instead. jax.interpreters.xla.backend_specific_translations _BackendSpecificTranslationsAdapter(None,...

The issue is in your error message: ``` Found cuBLAS version 120205, but JAX was built against version 120304, which is newer. ``` You'll need to install cuBLAS version 12.3...

`jax[cuda12_pip]` installs the correct CUDA sources in your Python `site_packages`. However, if you have other CUDA installations on your system, and your system is set up to load those other...

Hi - thanks for the request! We generally follow the NumPy API in `jax.numpy`, and as far as I can tell, numpy does not support a `where` argument for any...

`fill_diagonal` is somewhat different: it's *impossible* to implement that in JAX without changing the API to avoid in-place modification. Adding additional functionality to existing APIs is a qualitatively different discussion....

We have a related discussion here: https://jax.readthedocs.io/en/latest/jep/18137-numpy-scipy-scope.html, although it kind of takes as given that we won't implement things in `jax.numpy` or `jax.scipy` that are not in `numpy` or `scipy`...

Thanks for elaboration. I've seen this kind of idea come up before in other contexts (e.g. https://github.com/numpy/numpy/issues/24081) I personally don't feel strongly about this (it would be convenient in some...