Question

Saving Gradient in Backward Pass Google-JAX

score 529 · Answer 1 · 2025-01-05 16:26:08

529

Views

Saving Gradient in Backward Pass Google-JAX

Published on 05 January 2025 at 16:26

score 59 · Answer 2 · 2025-01-05 16:28:03

59

Views

Optimizing sampling with varying sample sizes in jax

Published on 05 January 2025 at 16:28

score 590 · Answer 3 · 2025-01-05 16:29:44

590

Views

Jax vmap, in_axes doesn't work if keyword argument is passed

Published on 05 January 2025 at 16:29

score 598 · Answer 4 · 2025-01-05 16:25:36

598

Views

How to use grad convolution in google-jax?

Published on 05 January 2025 at 16:25

score 179 · Answer 5 · 2025-01-05 16:32:25

179

Views

Profiling JAX code: What is redzone_checker and why does it take so much time?

Published on 05 January 2025 at 16:32

score 1.4k · Answer 6 · 2025-01-05 16:31:44

1.4k

Views

Is it possible to use objects with Google's Jax machine learning library

Published on 05 January 2025 at 16:31

score 540 · Answer 7 · 2025-01-05 16:25:35

540

Views

JAX/XLA slow compilation using conda

Published on 05 January 2025 at 16:25

score 386 · Answer 8 · 2025-01-05 16:32:52

386

Views

How to vectorize JAX functions using jit compilation and vmap auto-vectorization

Published on 05 January 2025 at 16:32

score 67 · Answer 9 · 2025-01-05 16:23:40

67

Views

Can jax.vmap() do a hstack()?

Published on 05 January 2025 at 16:23

score 87 · Answer 10 · 2025-01-05 16:27:10

87

Views

Does jax save the jaxpr of jit compiled functions?

Published on 05 January 2025 at 16:27

score 161 · Answer 11 · 2025-01-05 16:31:43

161

Views

jax.errors.UnexpectedTracerError only when using jax.debug.breakpoint()

Published on 05 January 2025 at 16:31

score 813 · Answer 12 · 2025-01-08 03:08:46

813

Views

JAX: Getting rid of zero-gradient

Published on 08 January 2025 at 03:08

score 655 · Answer 13 · 2025-01-05 16:32:25

655

Views

JAX grad function: why am I getting a list of zeros instead of the gradients?

Published on 05 January 2025 at 16:32

score 772 · Answer 14 · 2025-01-05 16:23:45

772

Views

Jax.lax.scan with arguments?

Published on 05 January 2025 at 16:23

score 2k · Answer 15 · 2025-01-05 16:25:32

2k

Views

Rewriting for loop with jax.lax.scan

Published on 05 January 2025 at 16:25

List Question

Saving Gradient in Backward Pass Google-JAX

Optimizing sampling with varying sample sizes in jax

Jax vmap, in_axes doesn't work if keyword argument is passed

How to use grad convolution in google-jax?

Profiling JAX code: What is redzone_checker and why does it take so much time?

Is it possible to use objects with Google's Jax machine learning library

JAX/XLA slow compilation using conda

How to vectorize JAX functions using jit compilation and vmap auto-vectorization

Can jax.vmap() do a hstack()?

Does jax save the jaxpr of jit compiled functions?

jax.errors.UnexpectedTracerError only when using jax.debug.breakpoint()

JAX: Getting rid of zero-gradient

JAX grad function: why am I getting a list of zeros instead of the gradients?

Jax.lax.scan with arguments?

Rewriting for loop with jax.lax.scan

Trending Questions

Popular # Hahtags

Popular Questions