HHälvä comments

Results 6 comments of


                                            HHälvä

vmap of cond's predicate results in select, leading to unexpected compute/memory use

I find this one of the biggest practical issues with jax -- `Vmap`+`jit` are great but in a lot of code this also necessitates use of `cond` with them, which...

vmap of cond's predicate results in select, leading to unexpected compute/memory use

Btw does `switch` suffer from this same problem when used with `vmap`?

vmap of cond's predicate results in select, leading to unexpected compute/memory use

That's what I thought -- might be helpful to document that for `switch` similarly to what's in `cond` doc.: e.g" However, when transformed with `vmap` to operate over a batch...

jax.scipy.linalg.eigh_tridiagonal() doesn't implement calculation of eigenvectors

This would be great to have!

jax.scipy.linalg.eigh_tridiagonal() doesn't implement calculation of eigenvectors

That sounds tricky. Shame, it makes it difficult to implement any of the popular GP inference engines that use eigendecomp of tridiagonal matrices as a way to compute log-determinants cheaply.

jax.scipy.linalg.eigh_tridiagonal() doesn't implement calculation of eigenvectors

> > a way to compute log-determinants cheaply. > > Perhaps I'm misunderstanding. Isn't `log det(M) = sum_i lg( eigh[i] )` so eigenvalues are sufficient? `eigh` is still expensive no?...