-
Hello! I was trying to implement some networks involving Exp, however it says that there is no stax.Exp().
I was running this in google colab CPU environment, with latest version 0.3.8.
But in the s…
-
I am working on empirical NTK calculations, and I feel the GPU memory usage is very high.
Two cases:
1. Input (28*28) => 1 Linear layer of 1024 channels => output 10 classes; Use 64 samples to cal…
-
In `Disentangling_Trainability_and_Generalization.ipynb`
```python
def mean_predictor(ntk_fn, x_train, y_train, x_test, reg=0.):
ktd = ntk_fn(x_test, x_train)
kdd = ntk_fn(x_train, x_train)
…
-
Hi,
I sincerely hope that my comments here will be purely seen in the spirit of constructive feedback.
I do not comment (not anymore!) when I see fragmentation and duplication of implementations…
-
Hi, would you mind explaining the difference between a general `Kernel` type and a `` type? Thanks ahead!
-
As can be seen in the example below, I expected that decorating a kernel function with `batch` doesn't change the return shape of the kernel function, i.e. both `res_5` and `res_10` below should have …
-
Hello! This is a very cool project!
I have a few questions related to using `jax-md` for condensed matter MD, in particular related to computing the stress. I'll give some context below, feel free …
-
Your work is really impressive. I visited your git to reproduce your result hoping that it could be the baseline for my research.
But there seem to be some issues here.
An error "'ShapedArray' obje…
-
Hi, thanks for this handy library. I have a basic question. I was trying to apply nngp on graph convolutional layers. For example,
```
# graph-nngp
init_fn, apply_fn, kernel_fn = stax.serial(
…
-
I wrote simple codes with **monte_carlo_kernel_fn** and **gradient_descent** modules, but
it raised an unidentifiable type error even though I've never manipulated any types in the code.
Basically, …