Roman Novak

Friendly ping - does the new `device_assignment` or `jax.device_put` address this by any chance? https://jax.readthedocs.io/en/latest/jax.html#module-jax I'm not sure what is the "XLA’s DeviceAssignment logic" and whether it allows to specify...

FYI, I have revisited the example below on: 1) **CPU**: einsum is slow AND wrong: https://colab.research.google.com/gist/romanngg/48fb8d4d3a3fb5da9be84d8d1fb862ad/einsum_is_wrong_and_slow_cpu.ipynb 2) **GPU**: einsum is slow: https://colab.research.google.com/gist/romanngg/dd1e2adbda90749f140012f1b9342353/einsum_is_slow_gpu.ipynb 3) **TPU**: einsum is OK! https://colab.research.google.com/gist/romanngg/635b467426bd9ead276cc6f9216ed03d/einsum_is_ok_tpu.ipynb Will file...

Haven't heard anything back yet

Yes, I was thinking exactly about the (R, +, *) -> ({True, False}, OR, AND) correspondence. Integer convolutions would be great too, but perhaps for my specific usecase I was...

Thanks James. Just to clarify, do you mean that on other platforms it will not let me save neither space, nor time compared to float32?

Related: #1273

yes that's another way to do it, we don't have a super-convenient function for it. We have https://neural-tangents.readthedocs.io/en/latest/monte_carlo.html for MC estimating kernels of any functions, but I assume you want...

Hi Vinith, yes I think this is correct, if something isn't working as expected let me know!

The discrepancy appears to be equal to the cost of convolving with the padded entries ``` 256 - 196 = 60 = 2 [multiply + add] * 2 [output filters]...

Mathematically I think we're doing what you've wrote, but we implement it with Cholesky factorization, so instead of ```python mean_predict.analytical result = kernel_cov.T.dot(Kff_inv).dot(train_ys) ``` we do something like ```python import...