Discussions

• 1 You must be logged in to vote

  🙏

  Example for flax.jax_utils.prefetch_to_device with TFDS dataset

  davidshen84 asked Apr 19, 2024 in Q&A · Closed · Unanswered

  

  

  2
• 1 You must be logged in to vote

  🙏

  How to init BatchNorm in the setup function

  davidshen84 asked Apr 23, 2024 in Q&A · Unanswered

  

  

  0
• 1 You must be logged in to vote

  🙏

  Training performance fluctuate

  davidshen84 asked Apr 18, 2024 in Q&A · Unanswered

  

  

  1
• 1 You must be logged in to vote

  🙏

  How to fix some layers for transfer learning?

  wztdream asked Dec 9, 2021 in Q&A · Answered

  

  

  

  

  11
• 1 You must be logged in to vote

  🙏

  Does flax/jax care about image channel order?

  davidshen84 asked Nov 22, 2023 in Q&A · Answered

  

  

  

  

  8
• 1 You must be logged in to vote

  🙏

  help with categorical encoding

  g00phy asked Apr 18, 2024 in Q&A · Closed · Unanswered

  

  

  0
• 1 You must be logged in to vote

  🙏

  fori_loop, scan and compatibility with linen.nn.Module

  Rasmuskh asked Apr 14, 2024 in Q&A · Unanswered

  

  

  0
• 2 You must be logged in to vote

  🙏

  How to load PyTorch checkpoints into JAX/Flax?

  marcvanzee asked Jan 18, 2021 in Q&A · Answered

  

  

  

  

  7
• 1 You must be logged in to vote

  🙏

  Access module function from outside

  braniii asked Apr 12, 2024 in Q&A · Unanswered

  

  

  0
• 1 You must be logged in to vote

  🙏

  Model Inspection: Computing Model FLOPs

  peterdavidfagan asked Apr 8, 2024 in Q&A · Closed · Answered

  

  

  1
• 1 You must be logged in to vote

  🙏

  Gradients leak through masked convolutions on GPU

  phlippe asked Aug 9, 2022 in Q&A · Answered

  

  

  

  

  6
• 1 You must be logged in to vote

  🙏

  more efficient dot_product_attention when having very long sequences with sparse mask

  TsingQAQ asked Apr 9, 2024 in Q&A · Unanswered

  

  

  0
• 3 You must be logged in to vote

  #️⃣

  MultiHeadDotProductAttention API

  cgarciae asked Dec 17, 2021 in General · Closed · Unanswered

  

  

  

  

  6
• 1 You must be logged in to vote

  🙏

  Best way to convert PyTorch recurrent model to Flax?

  Chulabhaya asked Apr 7, 2024 in Q&A · Unanswered

  

  

  0
• 1 You must be logged in to vote

  🙏

  What is the right way to scale up nn.MultiHeadDotProductAttention and nn.LayerNorm in multiple device?

  sw32-seo asked Apr 3, 2024 in Q&A · Unanswered

  

  

  0
• 1 You must be logged in to vote

  #️⃣

  nnx and jax.experimental.attrs

  ASEM000 asked Mar 31, 2024 in General · Answered

  

  

  2
• 1 You must be logged in to vote

  🙏

  Error when calling module tabulate involving WeightNorm

  DBraun asked Feb 19, 2024 in Q&A · Closed · Answered

  

  

  1
• 1 You must be logged in to vote

  🙏

  Model inference speed is slower than training speed

  davidshen84 asked Mar 24, 2024 in Q&A · Closed · Answered

  

  

  1
• 1 You must be logged in to vote

  🙏

  Vmap Flax + Huggingface Model

  pharringtonp19 asked Mar 21, 2024 in Q&A · Unanswered

  

  

  0
• 1 You must be logged in to vote

  🙏

  Using flax.nn.scan with modules that have callable containing kwargs

  Priority: P2 - no schedule Best effort response and resolution. We have no plan to work on this at the moment.

  peterdavidfagan asked Mar 15, 2024 in Q&A · Answered

  

  

  1
• 1 You must be logged in to vote

  🙏

  Preventing tabulate from truncating in Google Colab?

  edmuthiah asked Mar 13, 2024 in Q&A · Closed · Answered

  

  

  1
• 1 You must be logged in to vote

  🙏

  Does flax have an implementation of a full decoder-only model?

  epignatelli asked Feb 20, 2024 in Q&A · Unanswered

  

  

  

  

  3
• 3 You must be logged in to vote

  🙏

  using flax sow inside a function with scanning

  Priority: P2 - no schedule Best effort response and resolution. We have no plan to work on this at the moment.

  TalKachman asked Feb 28, 2024 in Q&A · Unanswered

  

  

  0
• 2 You must be logged in to vote

  🙏

  Lowering of normalization calls

  pratnali-aws asked Feb 21, 2024 in Q&A · Unanswered

  

  

  0
• 1 You must be logged in to vote

  🙏

  Remat (aka checkpoint) with tied variables from custom variable collections

  giganttheo asked Feb 1, 2024 in Q&A · Closed · Unanswered

  

  

  2