All pull requests

Pull requests

• Fix variance scaling with negative fan-in axes

  #908·Linxiushen opened Sep 27, 2026·

  ·
  1

• Preserve input shapes when reusing GroupNorm

  #907·sylvesterkaczmarek opened Sep 23, 2026·

  ·

• Honor negative axes when sizing BatchNorm parameters

  #906·sylvesterkaczmarek opened Sep 21, 2026·

  ·

• Preserve broadcast batch dimensions in MultiHeadAttention

  #905·sylvesterkaczmarek opened Sep 21, 2026·

  ·

• refactor(batch_norm): reduce test duplication and improve readability

  #890·pythongiant opened Jun 15, 2026·

  ·
  1

• Improve EMA documentation clarity and stability test coverage

  #888·savirpatil opened May 20, 2026·

  ·

• Avoid usage of deprecated jax.core APIs.

  #879·copybara-service[bot] opened Apr 15, 2026Bot·

  ·

• Add runtime deprecation warnings to one_hot and multinomial

  #865·t-suraj21 opened Jan 30, 2026·

  ·
  1

• Upgrade GitHub Actions for Node 24 compatibility

  #863·salmanmkc opened Jan 25, 2026·

  ·
  1

• Improve beginner explanations in Haiku basics notebook

  #861·MukarramBambot opened Jan 19, 2026·

  ·
  1

• fix(batch_norm): allow string cross_replica_axis for vmap support

  #859·aditya161205 opened Jan 8, 2026·

  ·
  2

  1

• [pmap] Remove jax_pmap_shmap_merge configuration.

  #850·copybara-service[bot] opened Dec 3, 2025Bot·

  ·

• Add Feedback Alignment implementation

  #848·hwilner opened Nov 1, 2025·

  ·
  1

  1

• Disable running tests on 3.1{0,1}.

  #839·copybara-service[bot] opened Aug 4, 2025Bot·

  ·

• Fix incorrect type specification for Haiku partition function.

  #826·copybara-service[bot] opened Jun 11, 2025Bot·

  ·
  1

• fixing avg_pool wrong shape issue

  #824·Vasist10 opened Jun 1, 2025·

  ·
  1

• Missing MaxPool2d

  #815·DEBADAS001KERNEL opened Apr 5, 2025·

  ·

• Avoid depending on JAX internals, which are about to change (see cl/681582933).

  #801·copybara-service[bot] opened Oct 23, 2024Bot·

  ·

• Update libraries to use JAX's limited (and ill-advised) trace-state-querying APIs rather than depending on JAX's deeper internals, which are about to change.

  #794·copybara-service[bot] opened Sep 23, 2024Bot·

  ·

• Fix packaging issue

  #753·ethanluoyc opened Nov 1, 2023Contributor·

  ·
  1

• Replace jax.linear_util -> jax.extend.linear_util to suppress DeprecationWarnings

  #739·copybara-service[bot] opened Oct 6, 2023Bot·

  ·

• Replace deprecated jax.linear_util.wrap_init with jax.extend.linear_util.wrap_init.

  #722·copybara-service[bot] opened Sep 1, 2023Bot·

  ·

• modify layer_stack transparency map

  #707·copybara-service[bot] opened Jul 27, 2023Bot·

  ·
  1

• MultiHeadAttention memory usage reduction via tiling

  #679·maxconway opened Jun 23, 2023·

  ·

• Fix setup.py packages to avoid installing "examples" and "docs"

  #672·natsukium opened Jun 19, 2023·

  ·