Yuxuan Jiang

It works, but I have to write an extra wrapper function to convert the wrapper type back, which feels redundant.

https://gist.github.com/jjyyxx/f64e28f6ccc37c24af9fd17649710b26 I created a demo that successfully saves a traced JAX function using pickle only. This is important to me because during quick research, I often make many small tweaks...

``` +-----------------------------------------------------------------------------------------+ | NVIDIA-SMI 550.54.15 Driver Version: 550.54.15 CUDA Version: 12.4 | |-----------------------------------------+------------------------+----------------------+ | GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC | | Fan Temp Perf Pwr:Usage/Cap...

@mdouze If I understand correctly, the `obj` is the squared-sum of all sample errors. But if I pass the `weights` parameter to `KMeans.train`, the objective matches neither the weighted (scikit-learn...

Thanks for your suggestion! Indeed, I found that `jax.vmap` works just OK before filing this issue. But I was worried about the documentation saying `hk.vmap` is `Equivalent to jax.vmap() with...

> warning: Linking two modules of different target triples: 'LLVMDialectModule' is 'nvptx64-nvidia-gpulibs' whereas '' is 'nvptx64-nvidia-cuda' Not using this repo, but came across this warning after upgrading to JAX 0.4.14...

@kno10 Thanks for your suggestion! I tried it out, but it does not quite work as expected. Considering k-means objective function, it could be expressed as ```math L(Z, A) =...

@kno10 Thanks for the suggestion. This indeed works, but only after small tweaks to the source code. When the dissimilarity matrix is symmetric, everything works fine; but when it turns...

I mean that in this case, the field annotations (e.g., in/out) and the method annotations (e.g., self.in_size/self.out_size) refer to the same concept but use very different naming styles. ```python class...

@patrick-kidger I went through the jaxtyping internals and found that implementing the feature I requested wasn't too difficult. I've created a prototype and would appreciate your feedback on it. While...