diff --git a/declearn/model/__init__.py b/declearn/model/__init__.py index 71b604e8ad37c380324cd623c2294259ffc54327..927518ed4911df828b15db23cd308becea8bf87f 100644 --- a/declearn/model/__init__.py +++ b/declearn/model/__init__.py @@ -42,10 +42,13 @@ The automatically-imported submodules implemented here are: Optional Submodules ------------------- The optional-dependency-based submodules that may be manually imported are: -* haiku: jax- and haiku-interfacing tools - - HaikuModel: Model to wrap a haiku-transformable model function. - - JaxNumpyVector: Vector for jax array data structures. +* [haiku][declearn.model.haiku]: + Jax- and Haiku-interfacing tools. + - [HaikuModel][declearn.model.haiku.HaikuModel]: + Model to wrap a haiku-transformable model function. + - [JaxNumpyVector][declearn.model.haiku.JaxNumpyVector]: + Vector for jax array data structures. * [tensorflow][declearn.model.tensorflow]: TensorFlow-interfacing tools - [TensorflowModel][declearn.model.tensorflow.TensorflowModel]: @@ -68,7 +71,7 @@ from . import api from . import sklearn OPTIONAL_MODULES = [ - "jax", + "haiku", "tensorflow", "torch", ]