# \#jax

**URL:** https://discourse.julialang.org/tag/jax/958.md

[Latest](https://discourse.julialang.org/latest.md) · [Categories](https://discourse.julialang.org/categories.md) · [Tags](https://discourse.julialang.org/tags.md)

---

## [Efficient automatic differentation for Julia version \`jax.scan\`?](https://discourse.julialang.org/t/efficient-automatic-differentation-for-julia-version-jax-scan/132853)

<div class="topic-metadata">

**Author:** [@roflmaostc](https://discourse.julialang.org/u/roflmaostc)\
**Replies:** 4\
**Last updated:** [October 3, 2025, 4:14pm UTC](https://discourse.julialang.org/t/efficient-automatic-differentation-for-julia-version-jax-scan/132853 "2025-10-03T16:14:03Z")

</div>

Hi, I am using in some JAX array code jax.lax.scan which implements something like this in my example: function jax\_lax\_scan(f, x; accumulator\_init) acc = accumulator\_init for i in axes(x, 1) acc = …

---

## [Julia (AcceleratedKernels) vs JAX time comparison](https://discourse.julialang.org/t/julia-acceleratedkernels-vs-jax-time-comparison/129781)

<div class="topic-metadata">

**Author:** [@mcp](https://discourse.julialang.org/u/mcp)\
**Replies:** 21\
**Last updated:** [June 11, 2025, 1:54pm UTC](https://discourse.julialang.org/t/julia-acceleratedkernels-vs-jax-time-comparison/129781 "2025-06-11T13:54:33Z")

</div>

Hi everyone! Following this topic, I’d like to understand how to write an efficient CPU/GPU agnostic code using AcceleratedKernels.jl. My problem is writing the joint PDF of two random variables, one of which is condit…

---

## [Understanding Tullio performances](https://discourse.julialang.org/t/understanding-tullio-performances/129760)

<div class="topic-metadata">

**Author:** [@mcp](https://discourse.julialang.org/u/mcp)\
**Replies:** 12\
**Last updated:** [June 10, 2025, 12:25pm UTC](https://discourse.julialang.org/t/understanding-tullio-performances/129760 "2025-06-10T12:25:57Z")

</div>

Hi everyone! I’m new to Julia and I’m trying to translate a JAX code in Julia using Tullio macro that makes the code simple and clean. However, I’m not able to get the same performances even with simple functions. Here…

---

## [\`Zygote.gradient\` is 54000 TIMES slower than \`jax.gradient\`](https://discourse.julialang.org/t/zygote-gradient-is-54000-times-slower-than-jax-gradient/125396)

<div class="topic-metadata">

**Author:** [@ForceBru](https://discourse.julialang.org/u/ForceBru)\
**Replies:** 80\
**Last updated:** [February 1, 2025, 9:25pm UTC](https://discourse.julialang.org/t/zygote-gradient-is-54000-times-slower-than-jax-gradient/125396 "2025-02-01T21:25:46Z")

</div>

I’m almost unable to compute ant gradients (EDIT 9 hrs later: meant to say “any gradients”, but got gradients for ants) in Zygote. I initially attempted to optimize my model using around a million parameters, but it got …

---

## [arXiv: "The State of Julia for Scientific Machine Learning" by Berman & Ginesin](https://discourse.julialang.org/t/arxiv-the-state-of-julia-for-scientific-machine-learning-by-berman-ginesin/121413)

<div class="topic-metadata">

**Author:** [@jonasaugust](https://discourse.julialang.org/u/jonasaugust)\
**Replies:** 29\
**Last updated:** [December 15, 2024, 7:01pm UTC](https://discourse.julialang.org/t/arxiv-the-state-of-julia-for-scientific-machine-learning-by-berman-ginesin/121413 "2024-12-15T19:01:57Z")

</div>

The State of Julia for Scientific Machine Learning A brief survey. Considers why Julia isn’t more popular vs. Python & Jax. Finds calling Julia from other languages difficult.

---

## [PythonCall spends a lot of time showing stuff for JAX](https://discourse.julialang.org/t/pythoncall-spends-a-lot-of-time-showing-stuff-for-jax/121346)

<div class="topic-metadata">

**Author:** [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Replies:** 3\
**Last updated:** [October 15, 2024, 8:59pm UTC](https://discourse.julialang.org/t/pythoncall-spends-a-lot-of-time-showing-stuff-for-jax/121346 "2024-10-15T20:59:43Z")

</div>

Hi everyone, especially @cjdoris! I’m trying to call JAX from Julia and I find that PythonCall.jl has a lot of overhead in a simple conversion case due to… printing? Has anyone run into the same issue? Setup code: usi…

---

## [Lux.jl vs Jax](https://discourse.julialang.org/t/lux-jl-vs-jax/119654)

<div class="topic-metadata">

**Author:** [@ijk](https://discourse.julialang.org/u/ijk)\
**Replies:** 6\
**Last updated:** [September 21, 2024, 6:12pm UTC](https://discourse.julialang.org/t/lux-jl-vs-jax/119654 "2024-09-21T18:12:02Z")

</div>

I’ve started using Julia this semester, and would like to reimplement some ML models from Jax to Lux.jl. I have three concerns: I noticed there aren’t constructs like vmap. My existing Jax code has quite a few vmap cal…

---

## [Zygote much slower than JAX for automatic differentiation of energy](https://discourse.julialang.org/t/zygote-much-slower-than-jax-for-automatic-differentiation-of-energy/114239)

<div class="topic-metadata">

**Author:** [@albertomercurio](https://discourse.julialang.org/u/albertomercurio)\
**Replies:** 22\
**Last updated:** [May 15, 2024, 6:31am UTC](https://discourse.julialang.org/t/zygote-much-slower-than-jax-for-automatic-differentiation-of-energy/114239 "2024-05-15T06:31:51Z")

</div>

Hello, I’m new in the neural network field, and I’m starting to study Neural Network Quantum States (for example ground state searching using neural networks). Usually, they are implemented in Python using Jax, but then…

---

## [Thoughts on JAX vs Julia](https://discourse.julialang.org/t/thoughts-on-jax-vs-julia/86463)

<div class="topic-metadata">

**Author:** [@jacobusmmsmit](https://discourse.julialang.org/u/jacobusmmsmit)\
**Replies:** 21\
**Last updated:** [October 25, 2022, 11:06am UTC](https://discourse.julialang.org/t/thoughts-on-jax-vs-julia/86463 "2022-10-25T11:06:40Z")

</div>

Hi everyone. I just posted to Forem about my thoughts on using JAX and Julia: I didn’t write it in the post, but I’d highly recommend everyone who hasn’t tried JAX yet to give it a go (and I would say the same to JAX …

---

## [Interesting quote about JAX by Viral Shah in an article in Business Insider](https://discourse.julialang.org/t/interesting-quote-about-jax-by-viral-shah-in-an-article-in-business-insider/82978)

<div class="topic-metadata">

**Author:** [@akvsak](https://discourse.julialang.org/u/akvsak)\
**Replies:** 0\
**Last updated:** [June 18, 2022, 5:34pm UTC](https://discourse.julialang.org/t/interesting-quote-about-jax-by-viral-shah-in-an-article-in-business-insider/82978 "2022-06-18T17:34:28Z")

</div>

“JAX is a feat of engineering,” said Viral Shah, creator of the Julia programming language that experts frequently compare to JAX. “I think of JAX as a separate programming language that happens to be instantiated throu…

---

## [\[ANN\] DLPack.jl - Share CPU and CUDA arrays between Julia and Python](https://discourse.julialang.org/t/ann-dlpack-jl-share-cpu-and-cuda-arrays-between-julia-and-python/76871)

<div class="topic-metadata">

**Author:** [@PabloZubieta](https://discourse.julialang.org/u/PabloZubieta)\
**Replies:** 0\
**Last updated:** [February 21, 2022, 11:01pm UTC](https://discourse.julialang.org/t/ann-dlpack-jl-share-cpu-and-cuda-arrays-between-julia-and-python/76871 "2022-02-21T23:01:29Z")

</div>

Hello, I’m pleased to announce the release of DLPack.jl. DLPack is a C API that has continuously been adopted as a the base protocol for exchanging tensor data structures between different python libraries including JAX,…