-
In the notebook "Neural Tangents Cookbook", there was a section to draw sample functions from the prior .
Is there a way to to draw sample functions from the posterior (conditioned on the training…
-
Hello,
I implemented a brutally simple infinite-width model, calling the kernel_fn with a batch of a single vector.
When I run this on CPU, I don't run into any exorbitant memory issues.
How…
-
I am aware that the default inference implemented is based on Mean squared error (MSE) loss. Is there an implemented example or a way to obtain aleatoric uncertainty instead (either homoscedestic or h…
-
The NN with erf function output activation can occassionally output way beyond the boundary [-1,1]:
```
from jax import random
from neural_tangents import stax
import neural_tangents as nt
impo…
-
For a model being used for classification with `k` classes, for `n` datapoints, the NTK should be of the size `nk` X `nk`. How would we get that with neural-tangents?
Currently, I'm able to get a `…
-
The LLVM compiler pass uses excessive amounts of memory for deep networks which are constructed like this
```
stax.serial([my_layer]*depth)
```
In fact, the compilation may eventually OOM.
The …
-
Hey!
I was curious about why you choose this as the batch size:
https://github.com/google/neural-tangents/blob/9f2ebc88905c46d60b7c4a9da25636924acc9d45/neural_tangents/utils/batch.py#L611
I'm n…
-
I have an empirical kernel function that works on a small batch of inputs.
I then decorated the kernel function with batch decorator like below, then it gives out of memory error. I would like to use…
-
Hello!
I want to implement my custom layer (Dense layer without bias term (y = Wx) and with custom weight initialization).
How can I do this? I've looked at https://github.com/google/neural-tang…
-
Hello,
When I try to compute the NTK of a model with an embedding layer, I get the following warning:
```
/usr/local/lib/python3.10/dist-packages/neural_tangents/_src/empirical.py:2215: UserWarni…