Discussions

• 1 You must be logged in to vote

  🙏

  Max pooling with reduce_window causes "ValueError: Linearization failed ..." only on some cases

  yippiez asked Oct 31, 2025 in Q&A · Unanswered

  

  

  0
• 4 You must be logged in to vote

  💡

  Lower latency associative scan option

  oliverdutton asked May 6, 2022 in Ideas · Unanswered

  

  

  2
• 1 You must be logged in to vote

  🙏

  Does jax.vmap understand not to repeat scalar-only-derived computations?

  jadball asked Oct 30, 2025 in Q&A · Answered

  

  

  3
• 1 You must be logged in to vote

  🙏

  Jax.sparse.linalg.eigsh like scipy eigsh?

  checc1 asked Oct 28, 2025 in Q&A · Unanswered

  

  

  3
• 1 You must be logged in to vote

  #️⃣

  jax.numpy dtypes deeply violate the principle of least surprise

  davidmarttila asked Oct 25, 2025 in General · Unanswered

  

  

  

  

  6
• 1 You must be logged in to vote

  🙏

  Multi-node training using sharding API

  neel04 asked Mar 3, 2024 in Q&A · Unanswered

  

  

  

  

  9
• 1 You must be logged in to vote

  🙏

  How to use traced indices to select elements in a string array?

  fishjojo asked Oct 24, 2025 in Q&A · Answered

  

  

  1
• 1 You must be logged in to vote

  🙏

  Getting JAX to use bfloat16 everywhere

  cool-RR asked Oct 22, 2025 in Q&A · Answered

  

  

  10
• 1 You must be logged in to vote

  🙏

  JAX-based library for simple language model training and RLHF?

  cool-RR asked Oct 18, 2025 in Q&A · Unanswered

  

  

  2
• 1 You must be logged in to vote

  🙏

  Status and workarounds for slow lax.while_loop on GPU?

  cgiovanetti asked Oct 21, 2025 in Q&A · Answered

  

  

  2
• 1 You must be logged in to vote

  🙏

  Why does reduce_sum_p takes out_sharding as an argument, but reduce_max_p and reduce_min_p do not?

  botev asked Oct 23, 2025 in Q&A · Unanswered

  

  

  4
• 1 You must be logged in to vote

  💡

  Add jax_config.abstract_mesh_context_manager to core.get_sharding cache context.

  botev asked Oct 23, 2025 in Ideas · Answered

  

  

  3
• 1 You must be logged in to vote

  🙏

  How does Jax return numpy arrays of bfloat16 dtype?

  botev asked Oct 23, 2025 in Q&A · Answered

  

  

  1
• 1 You must be logged in to vote

  💡

  Way to allow int64 without f64

  botev asked Oct 22, 2025 in Ideas · Unanswered

  

  

  

  

  8
• 1 You must be logged in to vote

  #️⃣

  To what extent is the sparse.BCOO API subject to change?

  BenjaminDAnjou asked Oct 22, 2025 in General · Answered

  

  

  2
• 1 You must be logged in to vote

  🙏

  Why bias isn't optional in pallas rmsnorm kernel?

  AakashKumarNain asked Oct 15, 2025 in Q&A · Unanswered

  

  

  2
• 1 You must be logged in to vote

  🙏

  Strange recompilations

  botev asked Oct 9, 2025 in Q&A · Unanswered

  

  

  9
• 2 You must be logged in to vote

  🙏

  Global way to automatically convert all floating-point computations to bfloat16

  carlosgmartin asked Jul 10, 2025 in Q&A · Unanswered

  

  

  

  

  2
• 1 You must be logged in to vote

  #️⃣

  The Jax primitive jax.lax.conv_transpose result does not match the result of torch.nn.functional.conv_transpose2d

  bug Something isn't working

  SanjithKumar2 asked Oct 13, 2025 in General · Answered

  

  

  

  

  12
• 2 You must be logged in to vote

  🙏

  vjp3 and sharding

  PhilipVinc asked Sep 19, 2025 in Q&A · Unanswered

  

  

  

  

  6
• 1 You must be logged in to vote

  #️⃣

  About jit of multi-host pmap not implemented

  matrix97317 asked Oct 17, 2025 in General · Unanswered

  

  

  0
• 2 You must be logged in to vote

  🙏

  Parallelization of JIT compilation?

  naquad asked Oct 16, 2025 in Q&A · Unanswered

  

  

  0
• 1 You must be logged in to vote

  🙏

  Parallel setup / Auto Sharding mode for general tensor operation?

  DiagRisker asked Oct 9, 2025 in Q&A · Unanswered

  

  

  1
• 1 You must be logged in to vote

  🙏

  can't import "sici" from jax.scipy.special

  kovnerd asked Oct 15, 2025 in Q&A · Unanswered

  

  

  1
• 1 You must be logged in to vote

  🙏

  Profiling large jax scripts

  YannBerthelot asked Oct 13, 2025 in Q&A · Unanswered

  

  

  1