jax

130 Questions
Newest Active Unanswered Frequent
Score
View
Card Compact
Using JAX with NetworkX
user_194472270
•
1
1
32
jax networkx
JAX: How to accumulate jnp.array on a condition in a jit function?
user_7426340
•
1
1
72
jax jit python
AttributeError: module 'jaxlib' has no attribute 'version'
user_73396240
•
2
1
255
jax conda python-3.x package python
is there a way to trace grads through self.put_variable method in flax?
user_121393490
•
1
2
53
flax jax python
Why does JAX + STAX model take more GPU memory than needed?
user_181623890
•
1
0
275
oom stax jax memory-profiling gpu
Modify an array from indexes contained in another array
user_98176770
•
5
1
89
jax numpy arrays python
How to improve this toy Jax optimizer code with while loops and saved history?
user_10785290
•
2
2
164
jax
Calculating the Hessian Vector Product of a Flax NN output wrt to the inputs
user_63694510
•
2
1
90
hessian flax jax automatic-differentiation
How to get keys for jax.tree_flatten object?
user_2872380
•
1
1
110
jax
JAX pmap with multi-core CPU
user_186499920
•
1
0
452
pmap jax multicore
Hot Questions