Hi everyone. I just posted to Forem about my thoughts on using JAX and Julia: [image] Thoughts on using JAX and Julia Shared attributes I consider both languages to be advanced tools in the realm of... I didn’t write it in the post, but I’d highly recommen…

True, I rewrote the second half without touching the first and you’re right it doesn’t work anymore.

While i agree that jax is more restrictive in terms of types vs. Julia (unless you do a lot of pytree stuff), i would say jax is less restrictive in terms of its automatic differentiation capabilities. There’s things jax can do with a single method call that julia cant even dream of. Perfect exampl…

Interesting. I’m not very familiar with Jax so please forgive any naivety on my behalf here. I’m rather surprised by your comment as I was under the impression that this is exactly the kind of thing Julia would be extremely good for. Edge case models that do not have a super optimized python workflo…

[image] Craig_Hamel: There’s things jax can do with a single method call that julia cant even dream of. what about Jax can’t do if-else what about Jax can’t even handle variable-lenth vector (i.e. a normal vector) what about Jax can’t handle while loop (any non-unrolled control florw)

Well for PINNs it’s not just hard, I dont think its doable currently (unless theres a new AD package I am unaware of). I’ve tried with zygote and forward and reverse diff and kept running into either errors trying to AD some AD code or zygote just hung all night trying to compile. With jax, you don…

Jax can handle conditional branching, maybe not native python if else but they have a functional version. Sure they can’t handle varibale array sizes in some situations but I dont think that’s a game changer. You could likely overcome this with some pytree games. Jax can handle a while loop you ju…

Thanks for the reply. I’ll definitely look some more into Jax.

[image] Craig_Hamel: There’s things jax can do with a single method call that julia cant even dream of. Perfect example is obysics informed neural networks where you need to reverse diff over a loss function defined with forward diffs. There’some work arounds in NeuralPDE.jl using finite differe…

[image] ChrisRackauckas: It’ll be “the” thing for Julia in about 2 years will Diffractor ever be useful before useless?

I wasn’t trying to oversell jax or julia. Personally i wish julia’s ecosystem was where it needed to be for my purposes but ill sacrifice O(N^D) run time for O(1) developer time vs. O(p) developer time where p is the number of julia AD packages.

Thoughts on JAX vs Julia

Community

jling August 29, 2022, 7:50pm 11

will Diffractor ever be useful before useless?

Topic		Replies	Views
What happened to XLA.jl Machine Learning	16	5306	March 12, 2023
What lessons could Julia's autodiff ecosystem learn from Stan's TinyGrad? Machine Learning	45	5434	September 26, 2025
Performance comparison - Flux.jl's Adam vs Jax's Adam Performance question , package , performance , flux	33	4372	November 3, 2022
State of machine learning in Julia Machine Learning	60	67688	August 26, 2022
Automatic Differentiation (AD) in Julia vs. Python (or PyTorch) Machine Learning autodiff	14	2907	January 16, 2025

Thoughts on JAX vs Julia

Related topics