Skip to content

Pull requests: google-deepmind/dm-haiku

New pull request New
24 Open 574 Closed
24 Open 574 Closed
Author
Filter by author
Loading
Label
Filter by label
Loading
Use alt + click/return to exclude labels
or ⇧ + click/return for logical OR
Projects
Filter by project
Loading
Milestones
Filter by milestone
Loading
Reviews
Filter by reviews
No reviews Review required Approved review Changes requested
Assignee
Filter by who’s assigned
Assigned to nobody Loading
Sort
Sort by
Newest Oldest Most commented Least commented Recently updated Least recently updated Best match
Most reactions
👍 👎 😄 🎉 😕 ❤️ 🚀 👀

Pull requests list

[pmap] Remove jax_pmap_shmap_merge configuration.
#850 opened Dec 3, 2025 by copybara-service bot Loading…
Add Feedback Alignment implementation
#848 opened Nov 1, 2025 by hwilner Loading…
1
1
Added tests for LayerNorm dtype and jit
#846 opened Sep 21, 2025 by edawite Loading…
Disable running tests on 3.1{0,1}.
#839 opened Aug 4, 2025 by copybara-service bot Loading…
Fix incorrect type specification for Haiku partition function.
#826 opened Jun 11, 2025 by copybara-service bot Loading…
1
fixing avg_pool wrong shape issue
#824 opened Jun 1, 2025 by Vasist10 Loading…
1
Missing MaxPool2d
#815 opened Apr 5, 2025 by DEBADAS001KERNEL Loading…
Avoid depending on JAX internals, which are about to change (see cl/681582933).
#801 opened Oct 23, 2024 by copybara-service bot Loading…
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 opened Sep 23, 2024 by copybara-service bot Loading…
Fix packaging issue
#753 opened Nov 1, 2023 by ethanluoyc Loading…
1
Replace jax.linear_util -> jax.extend.linear_util to suppress DeprecationWarnings
#739 opened Oct 6, 2023 by copybara-service bot Loading…
Replace deprecated jax.linear_util.wrap_init with jax.extend.linear_util.wrap_init.
#722 opened Sep 1, 2023 by copybara-service bot Loading…
modify layer_stack transparency map
#707 opened Jul 27, 2023 by copybara-service bot Loading…
1
MultiHeadAttention memory usage reduction via tiling
#679 opened Jun 23, 2023 by maxconway Loading…
Fix setup.py packages to avoid installing "examples" and "docs"
#672 opened Jun 19, 2023 by natsukium Loading…
haiku.base: handle both new-style and old-style PRNG keys
#653 opened May 22, 2023 by copybara-service bot Loading…
Internal change.
#548 opened Oct 26, 2022 by copybara-service bot Loading…
1
Internal change.
#547 opened Oct 25, 2022 by copybara-service bot Loading…
1
Add repeated rng calling to docs.
#543 opened Oct 18, 2022 by copybara-service bot Loading…
[JAX] Add an option to control the out-of-bounds indexing behavior of Haiku embedding layers.
#493 opened Aug 5, 2022 by copybara-service bot Loading…
Fix "output without random key" variable reference
#347 opened Mar 29, 2022 by almostimplemented Loading…
2
Feature #88: AdaptiveAvgPool2D Implementation
#327 opened Mar 5, 2022 by 0xrushi Loading…
1
Internal
#314 opened Feb 22, 2022 by copybara-service bot Loading…
Bump JAX requirements.
#272 opened Dec 10, 2021 by copybara-service bot Loading…
ProTip! Add no:assignee to see everything that’s not assigned.