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

Changed how hk is defined in _src submodules #609

Closed copybara-service[bot] closed 1 year ago

copybara-service[bot] commented 1 year ago

Changed how hk is defined in _src submodules

The majority of Haiku modules are implemented in terms of Haiku public API. This was previously done via an ad-hoc types.ModuleTypes object mirroring the structure of the public API. For example,

hk = types.ModuleType("haiku") hk.Module = module.Module

That worked okay at runtime, but was opaque to Python tooling, which cannot reason about modules defined in such way.

This change implements an alternative approach, where hk is a class. Although not particularly pretty, the new version is easy to reason about for both humans and tooling.