Haiku 0.0.16

What's Changed

• Fix inspect.stack() error caused by no_flax shim module. 58b4130.
• Various changes to remain compatible with the latest JAX release.

Full Changelog: v0.0.15...v0.0.16

Haiku 0.0.15

• Allow using Haiku without Flax in Python >=3.13 (#843).
• Various changes to remain compatible with the latest JAX release.

Full Changelog: v0.0.14...v0.0.15

Haiku 0.0.14

• Various changes to remain compatible with latest JAX release (e.g. #796, #806, #807, #814).

Haiku 0.0.13

• Various changes to remain compatible with latest JAX release (e.g. #779, #783, #789).
• Add shard_map flops estimate to jaxpr_info (#791).
• Correct an oversight in TruncatedNormal's documentation (#793).

Haiku 0.0.12

• hk.LayerStack transparency map now supports custom pytree leaves 9be7c9a.
• hk.LayerStack now supports layers using state 134763d.
• Calling repr(self) in module constructors is now supported (fixing #428).

Haiku 0.0.11

• hk.layer_stack now allows transparent application (no prefix on module names).
• hk.MultiHeadAttention allows bias initializer to be configured or biases to be removed.
• hk.DepthwiseConvND now supports dilation.
• hk.dropout supports broadcast_dims.
• hk.BatchApply avoids an unnecessary h2d copy during tracing.
• hk.experimental.profiler_name_scopes has been removed, these are on by default.
• Added hk.map mirroring jax.lax.map.

Haiku 0.0.10

• Added hk.mixed_precision.push_policy.
• Added hk.experimental.{get_params,get_initial_state,get_current_state}.
• Added hk.experimental.{maybe_get_rng_sequence_state,maybe_replace_rng_sequence_state}.
• hk.switch now supports multiple operands.
• hk.get_parameter now supports init=None.
• hk.MethodContext now includes orig_class.
• hk.GetterContext now includes lifted_prefix_name.
• hk.layer_stack now allows parameter reuse.
• Haiku is now compatible with jax.enable_custom_prng.
• TruncatedNormal now exports lower and upper bounds.
• Haiku init/apply functions now return dict rather than Mapping.
• hk.dropout now supports broadcast_dims.

Haiku 0.0.9

What's Changed

• Support vmap where in_axes is a list rather than a tuple in 307cf7d
• Pass pmap axis specs optionally to make_model_info in d0ba451
• Remove use of jax_experimental_name_stack flag in dbc0b1f
• Add param_axis argument to RMSNorm to allow setting scale param shape in a4998a0
• Add documentation and error messages for w_init and w_init_scale to avoid confusion in #541
• Fix hk.while_loop carrying state when reserving variable sizes of rng keys. by @copybara-service in #551
• Add ensemble example to hk.lift documentation. by @copybara-service in #556

Full Changelog: v0.0.8...v0.0.9

Haiku 0.0.8

• Added experimental.force_name.
• Added ability to simulate a method name in experimental.name_scope.
• Added a config option for PRNG key block size.
• Added unroll parameter to dynamic_unroll.
• Remove use of deprecated jax.tree_* functions.
• Many improvements to our examples.
• Improve error messages in vmap.
• Support jax_experimental_name_stack in jaxpr_info.
• transform_and_run now supports a map on PRNG keys.
• remat now uses the new JAX remat implementation.
• Scale parameter is now optional in RMSNorm.

Haiku 0.0.7

• Bug fix: modules with leading zeros (e.g. linear_007) are now correctly handled. 7632aff
• Breaking change: hk.vmap(..) now requires split_rng to be passed.
• Breaking change: hk.jit was removed from the public API.
• Feature: we always add profiler name scopes to Haiku modules with the latest version of JAX.
• Added a tutorial on parameter sharing.
• Added hk.ModuleProtocol and hk.SupportsCall.
• Added cross_replica_axis to VectorQuantiser.
• Added allow_reuse argument to hk.lift.
• Added fan_in_axes to VarianceScaling initialiser.
• Added hk.custom_setter(..) to intercept hk.set_state(..).
• Added hk.Deferred.
• Added hk.experimental.transparent_lift(..) and hk.experimental.transparent_lift_with_state(..).
• Added hk.experimental.fast_eval_shape(..).
• Added hk.experimental.current_name().
• Added hk.experimental.DO_NOT_STORE. 2a6c034
• Added config APIs.
• Added support for new jax.named_call implementation.
• The HAIKU_FLATMAPPING env var is no longer used.
• hk.dropout(..) now supports dynamic rate.
• hk.without_apply_rng(..) now supports multi transformed functions.