google-deepmind dm-haiku issues

google-deepmind / dm-haiku

JAX-based neural network library

https://dm-haiku.readthedocs.io

Apache License 2.0

2.91k stars 231 forks source link

issues

Newest

Newest Most commented Recently updated Oldest Least commented Least recently updated

WIP: Avoid pmap issue with linear layer in haiku.

#804 copybara-service[bot] opened 1 week ago
0
Avoid calling np.testing.assert_array_equal on random keys

#803 copybara-service[bot] closed 3 weeks ago
0
Avoid depending on JAX internals, which are about to change.

#802 copybara-service[bot] closed 3 weeks ago
0
Avoid depending on JAX internals, which are about to change (see cl/681582933).

#801 copybara-service[bot] opened 3 weeks ago
0
Require Python 3.10+ for Haiku (same as JAX) and modernize code to match.

#800 copybara-service[bot] closed 1 month ago
0
Use short aliases for common tree operations in Haiku.

#799 copybara-service[bot] closed 1 month ago
0
Cut 0.0.14.dev release.

#798 copybara-service[bot] closed 1 month ago
0
Cut 0.0.13 release.

#797 copybara-service[bot] closed 1 month ago
0
Make trace dispatch purely a function of context rather than a function of both context and data. This lets us delete a lot of machinery for managing data-dependent tracing: levels, sublevels, post_process_call, new_base_main, custom_bind and so on.

#796 copybara-service[bot] closed 2 weeks ago
0
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.

#795 copybara-service[bot] closed 1 month ago
0
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 1 month ago
0
Correct an oversight in `TruncatedNormal`'s documentation.

#793 copybara-service[bot] closed 2 months ago
0
[Haiku] Add shard_map flops estimate to the jaxpr_info.

#791 copybara-service[bot] closed 2 months ago
1
Remove uses of deprecated xb, xc, and xe abbreviation from jax.interpreters.xla

#789 copybara-service[bot] closed 3 months ago
0
Integrating vmap with BatchNorm

#788 reemabdelrazek30 opened 3 months ago
0
Relax error message in statful_tests

#787 copybara-service[bot] closed 4 months ago
0
hk.BatchNorm with jax.vmap

#786 reemabdelrazek30 opened 4 months ago
0
hk.switch does not work inside a hk.vmap function when hk.set_state is used

#785 rsmath opened 4 months ago
1
[JAX] Update users of jax.tree.map() to be more careful about how they handle Nones.

#784 copybara-service[bot] closed 5 months ago
0
make haiku work with upcoming JAX change to tree_map (being more careful about

#783 copybara-service[bot] closed 5 months ago
0
Make flax an optional dependency.

#782 copybara-service[bot] closed 5 months ago
0
Bump required JAX and TF versions to latest.

#781 copybara-service[bot] closed 6 months ago
0
Consider make flax an optional dependency

#780 jjyyxx closed 5 months ago
1
Replace deprecated `jax.tree_*` functions with `jax.tree.*`

#779 copybara-service[bot] closed 6 months ago
0
Wrapping the ```init``` function inside ```jax.jit```

#778 ksmdnl closed 2 days ago
1
Silence some pytype errors.

#777 copybara-service[bot] closed 7 months ago
0
Replace deprecated jax.random.shuffle with jax.random.permutation

#776 copybara-service[bot] closed 7 months ago
0
Exposing the experimental _split_transpose JAX scan parameter in Haiku.

#775 copybara-service[bot] closed 7 months ago
1
TypeError: 'type' object is not subscriptable

#774 chenyzh28 closed 7 months ago
4
If I am understanding correctly, one line of the model code of

#773 copybara-service[bot] closed 8 months ago
0
Fix some JAX key reuse issues in haiku tests

#772 copybara-service[bot] closed 8 months ago
0
Cut 0.0.12 release.

#771 copybara-service[bot] closed 8 months ago
0
Cut 0.0.13.dev release.

#770 copybara-service[bot] closed 8 months ago
0
Populate `_auto_repr` before running module constructor.

#769 copybara-service[bot] closed 8 months ago
0
Use tf-nightly for Haiku testing.

#768 copybara-service[bot] closed 9 months ago
0
Deactivate failsafe in hk.flax.lift to avoid false-positives.

#767 copybara-service[bot] closed 9 months ago
1
Train multiple hk.nets.MLP with one optimizer

#766 rsmath closed 7 months ago
2
modify layer_stack transparency map

#765 copybara-service[bot] closed 9 months ago
1
Fix _compute_fans when fan_in_axes is specified and len(shape) is 2.

#764 copybara-service[bot] closed 5 months ago
1
Haiku: fix jaxpr_info_test for jax v0.4.24

#763 copybara-service[bot] closed 10 months ago
0
Skip broken test.

#762 copybara-service[bot] closed 10 months ago
0
enable state with layer stack

#761 copybara-service[bot] closed 9 months ago
1
Internal change

#760 copybara-service[bot] closed 11 months ago
1
fix test not to depend on details of JAX shape error type/message

#759 copybara-service[bot] closed 1 year ago
0
Add a comment to force_name to avoid overuse.

#758 copybara-service[bot] closed 1 year ago
0
Cut 0.0.12.dev release.

#757 copybara-service[bot] closed 1 year ago
0
Cut 0.0.11 release.

#756 copybara-service[bot] closed 1 year ago
0
`haiku.experimental.flax` is not part of newest pip release

#755 simon-bachhuber closed 1 year ago
1
Allow for different parameters to be created when reuse=True.

#754 copybara-service[bot] closed 1 year ago
0
Fix packaging issue

#753 ethanluoyc opened 1 year ago
0