Skip to content

chore: remove dm-haiku support from numpyro.contrib.module#2150

Merged
juanitorduz merged 1 commit intopyro-ppl:masterfrom
theorashid:deprecate-haiku
Mar 16, 2026
Merged

chore: remove dm-haiku support from numpyro.contrib.module#2150
juanitorduz merged 1 commit intopyro-ppl:masterfrom
theorashid:deprecate-haiku

Conversation

@theorashid
Copy link
Copy Markdown
Contributor

This is for #2148.

Haiku is deprecated. Leaving flax linen there.

cc: @juanitorduz @thomaspinder

Copy link
Copy Markdown
Member

@fehiepsi fehiepsi left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Thanks @theorashid!

@fehiepsi
Copy link
Copy Markdown
Member

Could you sync to master @theorashid ?

     Remove `haiku_module` and `random_haiku_module` from contrib.module,
     the haiku ProdLDA example code, all haiku-related tests, docs, and the
     `dm-haiku` dev dependency. Also add `eqx_module` and `random_eqx_module`
     to `__all__` and replace `test_haiku_compatible` with `test_flax_compatible`
     in test_esc_proxies.
@juanitorduz
Copy link
Copy Markdown
Collaborator

Great! Thanks!

@juanitorduz juanitorduz merged commit 6b4b094 into pyro-ppl:master Mar 16, 2026
9 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants