Roman Novak

Thanks for the detailed repro! This touches on a bit of a fragile part of the library, which may take a while to figure out how fix properly (namely, we...

@SiuMath and @sschoenholz may answer better, but I can give some brief comments: 1) Re changing the depth, your observation is correct. $C*$ diagram shows the fixed point correlation, i.e....

At a glance, library usage seems good to me! Perhaps one way to figure this out is to establish a baseline using some other method (kernel, neural network, etc), to...

Note that in your example you are doing inference with an infinitely-wide neural network (`kernel_fn`), so the width doesn't matter in this case. Also, the plot does look like the...

Thanks a lot Jens, could you share a code sample that would fail before this change? I think we should support `n_feat > 1`, so I wonder if this points...

I'm afraid not, you would need to write your own `stax` layer, defining `init_fn`, `apply_fn`, and `kernel_fn`, e.g. as https://github.com/google/neural-tangents/blob/9f21e6e4f21a279ebbb033ff924e1ebc4723e077/neural_tangents/_src/stax/linear.py#L749 To what extent you'll be able to reuse your existing...

Thanks for the report, your code correct! This actually looks like two bugs on our side: 1) `store_on_device` argument isn't working, and the kernel is stored on the GPU (I'm...

Oh thanks for the error message, I realized what's actually failing is ```python fx_train_inf, fx_test_inf = predict_fn(fx_train_0=fx_train_0, fx_test_0=fx_test_0, k_test_train=k_test_train) ``` and not the kernel computation. Indeed 24Gb is not enough...

How much RAM do you have? Does it work (on CPU, after your modifications) if you use 36k points? I suspect you'd need at least ~64 Gb of RAM, but...

Indeed in examples and visualizations we have only used the diagonal entires (marginal variances), but in general outputs are described by a Gaussian process (GP) with non-zero non-diagonal covariance entries...