-
Right now, if you try to pickle a jaxpr, you are given an error.
```python
>>> import jax
>>> j = jax.make_jaxpr(f)(1)
>>> j
{ lambda ; a:i32[]. let b:i32[] = add a 1 in (b,) }
>>> import clo…
-
Hi My test run failed and showed the following error messages. Can you please give me a pointer to what went wrong? Thanks.
```
Running tests under Python 3.11.10: /alphafold3_venv/bin/python
[ R…
-
Hi, I have an extremely strange bug that is slightly hard to reproduce :)
I am trying to integrate brax with SAC and want to use an ensemble version of the algorithm. When I use the ensemble, the bra…
-
### Description
See the attached test case. Compiling a relatively simple function such as the below:
```py
def f(x, y):
t = x * y
for _ in range(n):
t += x * y
…
-
```python
import jax.numpy as jnp
jnp.array([None]*3)
# ==> TypeError: Unexpected input type for array:
# Yet:
np.array([None]*3)
# ==> array([None, None, None], dtype=object)
```
Is this…
-
The scikit-learn implementation of Gaussian Processes (GPs) scales as *O(n^3)* where *n* is the number of data points. Other packages such as [celerite](https://github.com/dfm/celerite) or [george](ht…
-
Is there any interest in adding a quasi-Newton based optimizer? I was thinking of porting over:
https://github.com/tensorflow/probability/blob/master/tensorflow_probability/python/optimizer/bfgs.py…
-
Hi @cgarciae -- treex is cool, in particular I really like the idea of using different classes for "variable kinds" and how it plays pretty nicely with typing (though I was a bit confused about why yo…
-
Hi!
I'm new to the JAX ecosystem, have used PyTorch and TensorFlow extensively for over 5 years.
My issue is that I can't get PyTorch data loading to work with jax/flax with `num_workers>0`.
F…
-
##### Description
It appears that the gradle open api plugin jax-rs-spec generator (java server-side api without an implementation) version 4.3.1 is not adding additional annotations on top of model …