romanngg
commented
11 months ago 1) Re up/downsampling, we have the ImageResize layer https://neural-tangents.readthedocs.io/en/latest/_autosummary/neural_tangents.stax.ImageResize.html#neural_tangents.stax.ImageResize - would it work for you?
2) It's a bit hard to diagnose non-PSDness, I may need to see the full code. Note that A won't be PSD unless x1 == x2, and if x1 == x2, you may want to try computing it as A = kernel_fn(x1, None, 'ntk'), which can be a bit more numerically stable.
Lmk if this help!
HI, Thank you for the nice work. I am trying to do kernel ridge regression with ntk. to do that i need down and upsampling layers. the following code i am using.
https://gist.github.com/jaya-pret/7eeb6a7edb26d4e3ea54f027ad1a2d2a
now for a given input x1 = (n1, h, w, c) how i can add upsampling and downsampling in this . and the ntk is not positive definite, the initailisation is random.uniform for x1. i have also tried jnp.ones with low value but then the kernel value becomes all same! is there any layers i can add so that the kernel value changes for ntk and becomes positive definite, or it depends on the data type. thank you.