Discussions

• 1 You must be logged in to vote

  🙏

  How to add barrier into a jitted function?

  knightXun asked Jul 30, 2024 in Q&A · Unanswered

  

  

  0
• 1 You must be logged in to vote

  🙏

  Trace leak with vmap

  All4RAIN asked Jul 29, 2024 in Q&A · Unanswered

  

  

  2
• 1 You must be logged in to vote

  💡

  Grid Sample

  josauder asked Apr 2, 2024 in Ideas · Unanswered

  

  

  

  

  6
• 1 You must be logged in to vote

  🙏

  Enforcing single-thread behaviour from jax

  Kamuish asked Jul 29, 2024 in Q&A · Unanswered

  

  

  0
• 1 You must be logged in to vote

  🙏

  Disabling TF32 in JAX

  Honza9723 asked Jul 28, 2024 in Q&A · Answered

  

  

  2
• 1 You must be logged in to vote

  🙏

  Which XlaOps are used for random number generation in JAX?

  Ebanflo42 asked Jul 29, 2024 in Q&A · Answered

  

  

  1
• 1 You must be logged in to vote

  🙏

  Tracing into function with specific custom_jvp

  SNMS95 asked Jul 26, 2024 in Q&A · Unanswered

  

  

  4
• 1 You must be logged in to vote

  🙏

  Resolving jaxlib conflicts with CUDA11.2

  joshdorrington asked Jul 24, 2024 in Q&A · Unanswered

  

  

  1
• 1 You must be logged in to vote

  🙏

  Memory Optimization with Jit

  mj023 asked Jul 25, 2024 in Q&A · Unanswered

  

  

  5
• 2 You must be logged in to vote

  🙏

  Should I be using jax.pmap?

  skushnir123 asked Jul 25, 2024 in Q&A · Closed · Answered

  

  

  1
• 1 You must be logged in to vote

  🙏

  Tricky NonConcreteBooleanIndexError

  harryjulian asked Jul 24, 2024 in Q&A · Answered

  

  

  6
• 2 You must be logged in to vote

  🙏

  Jax doesn't default to cpu when cuda is not available

  thomasgaudelet asked Jul 24, 2024 in Q&A · Unanswered

  

  

  0
• 1 You must be logged in to vote

  #️⃣

  Each iteration of lax.scan requires a kernel launch on GPU backends. Can this be resolved?

  carlosgmartin asked Jul 24, 2024 in General · Unanswered

  

  

  0
• 2 You must be logged in to vote

  #️⃣

  Tensor function optimization with Jax

  MRiabov asked Jul 16, 2024 in General · Unanswered

  

  

  1
• 1 You must be logged in to vote

  🙏

  Implementing padding inside jit function

  tchatow asked Jul 15, 2024 in Q&A · Answered

  

  

  3
• 1 You must be logged in to vote

  💡

  Jax Scipy RBFInterpolator?

  FluffyCodeMonster asked Jul 17, 2024 in Ideas · Unanswered

  

  

  1
• 1 You must be logged in to vote

  🙏

  Data Parallelization in class based methods

  sanepunk asked Jul 21, 2024 in Q&A · Unanswered

  

  

  0
• 2 You must be logged in to vote

  🙏

  parallelization in jax when the data are more than the number of devices/cores

  yyang97 asked Jan 21, 2024 in Q&A · Answered

  

  

  

  

  2
• 1 You must be logged in to vote

  🙏

  Train fit using pytrees

  saisua asked Jul 12, 2024 in Q&A · Unanswered

  

  

  

  

  4
• 1 You must be logged in to vote

  🙏

  [pallas]Explanation for scratch_shapes for tpu

  Lime-Cakes asked Jul 15, 2024 in Q&A · Answered

  

  

  3
• 2 You must be logged in to vote

  #️⃣

  Plans with jax.experimental.export

  dionhaefner asked Feb 9, 2024 in General · Answered

  

  

  

  

  6
• 1 You must be logged in to vote

  🙏

  Computing the leading singular vector of a matrix without computing the full SVD - analog of scipy.sparse.linalg.svds

  tonysf asked Jun 29, 2024 in Q&A · Unanswered

  

  

  

  

  2
• 1 You must be logged in to vote

  #️⃣

  What replaced as_serialized_hlo_module_proto() following the deprecation of jax.xla_computation?

  joaospinto asked Jul 4, 2024 in General · Closed · Answered

  

  

  9
• 1 You must be logged in to vote

  #️⃣

  Calling pre-compiled JAX code from C++

  joaospinto asked Jun 28, 2024 in General · Closed · Unanswered

  

  

  7
• 1 You must be logged in to vote

  🙏

  Unrecognisable XLA flags in JAX built on TPU with custom XLA

  awav asked Jan 20, 2022 in Q&A · Unanswered

  

  

  

  

  2