---

name: flax-interop description: "Use Haiku's optional experimental Flax interop APIs to embed Haiku modules in Flax or lift Flax modules inside Haiku." disable-model-invocation: true metadata: disco-role: operating license: Apache 2.0

Flax Interop

Use this sub-skill when a task specifically mixes Haiku with Flax and needs haiku.experimental.flax behavior. Treat the APIs here as optional, experimental interop helpers: they are useful for preserving existing Haiku code inside a Flax program, or for calling a Flax linen.Module inside a Haiku transform.

Route by direction

• Haiku inside Flax: use hk.experimental.flax.Module(transformed) for an hk.transform/hk.transform_with_state result, or hk.experimental.flax.Module.create(hk.ModuleClass, *args, **kwargs) for a Haiku module class.
• Flax inside Haiku: call hk.experimental.flax.lift(flax_module, name="...") inside an outer hk.transform or hk.transform_with_state function.
• Variable inspection/conversion: use hk.experimental.flax.flatten_flax_to_haiku(collection) on one Flax collection such as variables["params"], not on the whole variables dictionary.
• Generic Flax work: route out of this sub-skill for ordinary Flax training loops, non-Haiku Flax APIs, or broad migration planning.

Read or run these bundled files

• Read references/api-reference.md for a compact API table, signature shapes, collection mapping rules, and state/RNG semantics.
• Read references/flax-interop.md for copyable recipes that choose between Haiku-in-Flax and Flax-in-Haiku workflows.
• Read references/troubleshooting.md when optional flax imports fail, variable trees look surprising, state does not update, RNG streams fail, or versions are incompatible.
• Run scripts/haiku_flax_smoke.py to verify the optional Flax dependency and a Module.create(hk.Linear, ...) init/apply smoke path on synthetic data.

Minimal working checks

1. Confirm the dependency path first: python scripts/haiku_flax_smoke.py --mode module-create.
2. For Haiku-in-Flax, inspect variables.keys() after mod.init(...); expect Flax collections such as "params" and, for stateful Haiku, "state".
3. For Flax-in-Haiku, inspect the outer Haiku params and state returned by init; lifted Flax params appear under the explicit name, and non-param Flax collections become Haiku state prefixes.
4. If exact initialized values differ between a direct Haiku init path and a Flax-wrapped init path, treat that as expected RNG-splitting behavior unless the same variables are being shared explicitly.