LICENSE
MANIFEST.in
README.md
requirements-flax.txt
requirements-jax.txt
requirements-test.txt
requirements.txt
setup.py
dm_haiku.egg-info/PKG-INFO
dm_haiku.egg-info/SOURCES.txt
dm_haiku.egg-info/dependency_links.txt
dm_haiku.egg-info/not-zip-safe
dm_haiku.egg-info/requires.txt
dm_haiku.egg-info/top_level.txt
docs/conf.py
docs/ext/coverage_check.py
examples/imagenet/__init__.py
examples/imagenet/dataset.py
examples/imagenet/train.py
examples/impala/__init__.py
examples/impala/actor.py
examples/impala/actor_test.py
examples/impala/agent.py
examples/impala/haiku_nets.py
examples/impala/learner.py
examples/impala/learner_test.py
examples/impala/run_catch.py
examples/impala/util.py
examples/rnn/__init__.py
examples/rnn/dataset.py
examples/rnn/train.py
examples/transformer/__init__.py
examples/transformer/dataset.py
examples/transformer/model.py
examples/transformer/train.py
haiku/__init__.py
haiku/config.py
haiku/data_structures.py
haiku/initializers.py
haiku/mixed_precision.py
haiku/nets.py
haiku/pad.py
haiku/py.typed
haiku/testing.py
haiku/typing.py
haiku/_src/__init__.py
haiku/_src/analytics.py
haiku/_src/attention.py
haiku/_src/attention_test.py
haiku/_src/base.py
haiku/_src/base_test.py
haiku/_src/basic.py
haiku/_src/basic_test.py
haiku/_src/batch_norm.py
haiku/_src/batch_norm_test.py
haiku/_src/bias.py
haiku/_src/bias_test.py
haiku/_src/config.py
haiku/_src/config_test.py
haiku/_src/conv.py
haiku/_src/conv_test.py
haiku/_src/data_structures.py
haiku/_src/data_structures_test.py
haiku/_src/deferred.py
haiku/_src/deferred_test.py
haiku/_src/depthwise_conv.py
haiku/_src/depthwise_conv_test.py
haiku/_src/dot.py
haiku/_src/dot_test.py
haiku/_src/embed.py
haiku/_src/embed_test.py
haiku/_src/eval_shape.py
haiku/_src/eval_shape_test.py
haiku/_src/filtering.py
haiku/_src/filtering_test.py
haiku/_src/group_norm.py
haiku/_src/group_norm_test.py
haiku/_src/initializers.py
haiku/_src/initializers_test.py
haiku/_src/jaxpr_info.py
haiku/_src/jaxpr_info_test.py
haiku/_src/layer_norm.py
haiku/_src/layer_norm_test.py
haiku/_src/layer_stack.py
haiku/_src/layer_stack_test.py
haiku/_src/lift.py
haiku/_src/lift_test.py
haiku/_src/mixed_precision.py
haiku/_src/mixed_precision_test.py
haiku/_src/module.py
haiku/_src/module_test.py
haiku/_src/moving_averages.py
haiku/_src/moving_averages_test.py
haiku/_src/multi_transform.py
haiku/_src/multi_transform_test.py
haiku/_src/no_flax.py
haiku/_src/no_flax_test.py
haiku/_src/pad.py
haiku/_src/pad_test.py
haiku/_src/pool.py
haiku/_src/pool_test.py
haiku/_src/random.py
haiku/_src/random_test.py
haiku/_src/recurrent.py
haiku/_src/recurrent_test.py
haiku/_src/reshape.py
haiku/_src/reshape_test.py
haiku/_src/rms_norm.py
haiku/_src/rms_norm_test.py
haiku/_src/spectral_norm.py
haiku/_src/spectral_norm_test.py
haiku/_src/stateful.py
haiku/_src/stateful_test.py
haiku/_src/summarise.py
haiku/_src/summarise_test.py
haiku/_src/test_utils.py
haiku/_src/test_utils_test.py
haiku/_src/transform.py
haiku/_src/transform_test.py
haiku/_src/typing.py
haiku/_src/typing_test.py
haiku/_src/utils.py
haiku/_src/utils_test.py
haiku/_src/flax/flax_module.py
haiku/_src/flax/flax_module_test.py
haiku/_src/flax/transform_flax.py
haiku/_src/flax/transform_flax_test.py
haiku/_src/flax/utils.py
haiku/_src/flax/utils_test.py
haiku/_src/integration/__init__.py
haiku/_src/integration/check_tracer_leaks_test.py
haiku/_src/integration/checkpoint_generate.py
haiku/_src/integration/checkpoint_test.py
haiku/_src/integration/checkpoint_utils.py
haiku/_src/integration/common.py
haiku/_src/integration/descriptors.py
haiku/_src/integration/descriptors_test.py
haiku/_src/integration/doctest_test.py
haiku/_src/integration/float64_test.py
haiku/_src/integration/half_test.py
haiku/_src/integration/hk_transforms_test.py
haiku/_src/integration/jax2tf_test.py
haiku/_src/integration/jax_transforms_test.py
haiku/_src/integration/jaxpr_info_test.py
haiku/_src/integration/numpy_inputs_test.py
haiku/_src/integration/rank_promotion_test.py
haiku/_src/integration/shim_hk_test.py
haiku/_src/integration/summarise_test.py
haiku/_src/integration/to_dot_test.py
haiku/_src/integration/typing_test.py
haiku/_src/nets/__init__.py
haiku/_src/nets/mlp.py
haiku/_src/nets/mlp_test.py
haiku/_src/nets/mobilenetv1.py
haiku/_src/nets/mobilenetv1_test.py
haiku/_src/nets/resnet.py
haiku/_src/nets/resnet_test.py
haiku/_src/nets/vqvae.py
haiku/_src/nets/vqvae_test.py
haiku/benchmarks/eval_shape.py
haiku/benchmarks/init.py
haiku/experimental/__init__.py
haiku/experimental/flax.py
haiku/experimental/jaxpr_info.py