jax

130 Questions
Newest Active Unanswered Frequent
Score
View
Card Compact
JAX: avoid just-in-time recompilation for a function evaluated with a varying number of elements along one axis
user_127697830
•
3
1
444
jax jit python-3.x python
Jax - Debugging NaN-values
user_142947500
•
4
1
976
recurrent-neural-network jax nan python
How to compute the joint probability density function from a joint cumulative density function in Jax?
user_172782140
•
1
1
133
probability-distribution probability-density jax automatic-differentiation python
AttributeError: module 'jaxlib.xla_extension' has no attribute 'PmapFunction'
user_171711410
•
2
1
3400
jax nlp python
GPU and `jax` performance mysteries
user_29772560
•
1
1
69
jax numpy python performance
Why does my JAX + Haiku code don't run on GPU?
user_173825980
•
1
2
1182
jax google-colaboratory gpu python
Vectorise nested vmap
user_44513150
•
2
1
220
jax numpy python
Turn a tf.data.Dataset to a jax.numpy iterator
user_172894630
•
4
2
692
jax tensorflow numpy-ndarray python
Jax/Flax (very) slow RNN-forward-pass compared to pyTorch?
user_142947500
•
3
1
810
recurrent-neural-network jax pytorch python performance
how to do curve fitting using google jax?
user_13191280
•
1
1
234
jax curve-fitting python
Hot Questions