Module v3 (#92)

poets-ai · Nov 15, 2020 · 19ec87b · 19ec87b
1 parent f47aa1b
commit 19ec87b
Show file tree

Hide file tree

Showing 193 changed files with 4,396 additions and 3,183 deletions.
diff --git a/.github/workflows/ci_test.yml b/.github/workflows/ci_test.yml
@@ -41,11 +41,11 @@ jobs:
       - name: Install Poetry
         uses: dschep/install-poetry-action@v1.2
         with:
-          version: 1.0.10
+          version: 1.1.4
 
       - name: Install Dependencies
         run: |
-          poetry export --dev -f requirements.txt > requirements.txt
+          poetry export --without-hashes --dev -f requirements.txt > requirements.txt
           pip install -r requirements.txt
           pip install jax
 

diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
@@ -2,7 +2,7 @@
 This is a short guide on how to start contributing to Elegy along with some best practices for the project.
 
 ## Setup
-We use `poetry` so the easiest way to setup a development environment is run
+We use `poetry >= 1.1.4`, the easiest way to setup a development environment is run:
 
 ```bash
 poetry config virtualenvs.in-project true --local

diff --git a/README.md b/README.md
@@ -10,9 +10,9 @@
 
 -----------------
 
-_Elegy is a Neural Networks framework based on Jax inspired by Keras and Haiku._  
+_Elegy is a Neural Networks framework based on Jax inspired by Keras._  
 
-Elegy implements the Keras API but makes changes to play better with Jax and gives more flexibility around [losses and metrics](https://poets-ai.github.io/elegy/guides/modules-losses-metrics/), it also ports Haiku's excellent [module system](https://poets-ai.github.io/elegy/guides/module-system/) and makes it easier to use. Elegy is in an early stage, feel free to send us your feedback!
+Elegy implements the Keras API but makes changes to play better with Jax and gives more flexibility around [losses and metrics](https://poets-ai.github.io/elegy/guides/modules-losses-metrics/) and excellent [module system](https://poets-ai.github.io/elegy/guides/module-system/) that makes it super easy to use. Elegy is in an early stage, feel free to send us your feedback!
 
 #### Main Features
 
@@ -108,8 +108,8 @@ class Linear(elegy.Module):
         self.units = units
 
     def call(self, x):
-        w = elegy.get_parameter("w", [x.shape[-1], self.units], initializer=jnp.ones)
-        b = elegy.get_parameter("b", [self.units], initializer=jnp.ones)
+        w = self.add_parameter("w", [x.shape[-1], self.units], initializer=jnp.ones)
+        b = self.add_parameter("b", [self.units], initializer=jnp.ones)
 
         return jnp.dot(x, w) + b
 ```

diff --git a/docs/api/Loss.md b/docs/api/Loss.md
@@ -1,5 +1,6 @@
 
 # elegy.Loss
+
 ::: elegy.losses.loss.Loss
     selection:
         inherited_members: true

diff --git a/docs/api/Metric.md b/docs/api/Metric.md
@@ -1,5 +1,6 @@
 
 # elegy.Metric
+
 ::: elegy.metrics.metric.Metric
     selection:
         inherited_members: true

diff --git a/docs/api/Model.md b/docs/api/Model.md
@@ -1,6 +1,7 @@
 
 # elegy.Model
-::: elegy.model.Model
+
+::: elegy.model.model.Model
     selection:
         inherited_members: true
         members:
@@ -17,6 +18,5 @@
             - train_on_batch
             - full_state
             - parameters
-            - seed
             - states
 
diff --git a/docs/api/Module.md b/docs/api/Module.md
@@ -1,17 +1,16 @@
 
 # elegy.Module
+
 ::: elegy.module.Module
     selection:
         inherited_members: true
         members:
             - __init__
             - call
-            - init
-            - apply
-            - reset
+            - add_parameter
             - get_parameters
             - set_parameters
-            - get_states
-            - set_states
-            - submodules
+            - reset
+            - init
+            - initialized
 
diff --git a/docs/api/PRNGSequence.md b/docs/api/PRNGSequence.md
diff --git a/docs/api/RNG.md b/docs/api/RNG.md
@@ -0,0 +1,6 @@
+
+# elegy.RNG
+
+::: elegy.random.RNG
+    selection:
+        inherited_members: true
diff --git a/docs/api/add_loss.md b/docs/api/add_loss.md
@@ -1,5 +1,6 @@
 
 # elegy.add_loss
-::: elegy.hooks.add_loss
+
+::: elegy.module.add_loss
     selection:
         inherited_members: true
diff --git a/docs/api/add_metric.md b/docs/api/add_metric.md
@@ -1,5 +1,6 @@
 
 # elegy.add_metric
-::: elegy.hooks.add_metric
+
+::: elegy.module.add_metric
     selection:
         inherited_members: true
diff --git a/docs/api/add_summary.md b/docs/api/add_summary.md
@@ -1,5 +1,6 @@
 
 # elegy.add_summary
-::: elegy.hooks.add_summary
+
+::: elegy.module.add_summary
     selection:
         inherited_members: true
diff --git a/docs/api/callbacks/CSVLogger.md b/docs/api/callbacks/CSVLogger.md
@@ -1,5 +1,6 @@
 
 # elegy.callbacks.CSVLogger
+
 ::: elegy.callbacks.csv_logger.CSVLogger
     selection:
         inherited_members: true

diff --git a/docs/api/callbacks/Callback.md b/docs/api/callbacks/Callback.md
@@ -1,5 +1,6 @@
 
 # elegy.callbacks.Callback
+
 ::: elegy.callbacks.callback.Callback
     selection:
         inherited_members: true

diff --git a/docs/api/callbacks/CallbackList.md b/docs/api/callbacks/CallbackList.md
@@ -1,5 +1,6 @@
 
 # elegy.callbacks.CallbackList
+
 ::: elegy.callbacks.callback_list.CallbackList
     selection:
         inherited_members: true
diff --git a/docs/api/callbacks/EarlyStopping.md b/docs/api/callbacks/EarlyStopping.md
@@ -1,5 +1,6 @@
 
 # elegy.callbacks.EarlyStopping
+
 ::: elegy.callbacks.early_stopping.EarlyStopping
     selection:
         inherited_members: true

diff --git a/docs/api/callbacks/History.md b/docs/api/callbacks/History.md
@@ -1,5 +1,6 @@
 
 # elegy.callbacks.History
+
 ::: elegy.callbacks.history.History
     selection:
         inherited_members: true

diff --git a/docs/api/callbacks/LambdaCallback.md b/docs/api/callbacks/LambdaCallback.md
@@ -1,5 +1,6 @@
 
 # elegy.callbacks.LambdaCallback
+
 ::: elegy.callbacks.lambda_callback.LambdaCallback
     selection:
         inherited_members: true

diff --git a/docs/api/callbacks/ModelCheckpoint.md b/docs/api/callbacks/ModelCheckpoint.md
@@ -1,5 +1,6 @@
 
 # elegy.callbacks.ModelCheckpoint
+
 ::: elegy.callbacks.model_checkpoint.ModelCheckpoint
     selection:
         inherited_members: true

diff --git a/docs/api/callbacks/RemoteMonitor.md b/docs/api/callbacks/RemoteMonitor.md
@@ -1,5 +1,6 @@
 
 # elegy.callbacks.RemoteMonitor
+
 ::: elegy.callbacks.remote_monitor.RemoteMonitor
     selection:
         inherited_members: true

diff --git a/docs/api/callbacks/TensorBoard.md b/docs/api/callbacks/TensorBoard.md
@@ -1,5 +1,6 @@
 
 # elegy.callbacks.TensorBoard
+
 ::: elegy.callbacks.tensorboard.TensorBoard
     selection:
         inherited_members: true

diff --git a/docs/api/callbacks/TerminateOnNaN.md b/docs/api/callbacks/TerminateOnNaN.md
@@ -1,5 +1,6 @@
 
 # elegy.callbacks.TerminateOnNaN
+
 ::: elegy.callbacks.terminate_nan.TerminateOnNaN
     selection:
         inherited_members: true

diff --git a/docs/api/get_dynamic_context.md b/docs/api/get_dynamic_context.md
@@ -0,0 +1,6 @@
+
+# elegy.get_dynamic_context
+
+::: elegy.module.get_dynamic_context
+    selection:
+        inherited_members: true
diff --git a/docs/api/get_losses.md b/docs/api/get_losses.md
@@ -0,0 +1,6 @@
+
+# elegy.get_losses
+
+::: elegy.module.get_losses
+    selection:
+        inherited_members: true
diff --git a/docs/api/get_metrics.md b/docs/api/get_metrics.md
@@ -0,0 +1,6 @@
+
+# elegy.get_metrics
+
+::: elegy.module.get_metrics
+    selection:
+        inherited_members: true
diff --git a/docs/api/get_parameter.md b/docs/api/get_parameter.md
diff --git a/docs/api/get_rng.md b/docs/api/get_rng.md
@@ -0,0 +1,6 @@
+
+# elegy.get_rng
+
+::: elegy.module.get_rng
+    selection:
+        inherited_members: true
diff --git a/docs/api/get_state.md b/docs/api/get_state.md
diff --git a/docs/api/get_static_context.md b/docs/api/get_static_context.md
@@ -0,0 +1,6 @@
+
+# elegy.get_static_context
+
+::: elegy.module.get_static_context
+    selection:
+        inherited_members: true
diff --git a/docs/api/get_summaries.md b/docs/api/get_summaries.md
@@ -0,0 +1,6 @@
+
+# elegy.get_summaries
+
+::: elegy.module.get_summaries
+    selection:
+        inherited_members: true
diff --git a/docs/api/hooks/add_loss.md b/docs/api/hooks/add_loss.md
diff --git a/docs/api/hooks/add_metric.md b/docs/api/hooks/add_metric.md
diff --git a/docs/api/hooks/add_summary.md b/docs/api/hooks/add_summary.md
diff --git a/docs/api/hooks/get_parameter.md b/docs/api/hooks/get_parameter.md
diff --git a/docs/api/hooks/get_state.md b/docs/api/hooks/get_state.md
diff --git a/docs/api/hooks/next_rng_key.md b/docs/api/hooks/next_rng_key.md
diff --git a/docs/api/hooks/set_state.md b/docs/api/hooks/set_state.md
diff --git a/docs/api/hooks_context.md b/docs/api/hooks_context.md
@@ -0,0 +1,6 @@
+
+# elegy.hooks_context
+
+::: elegy.module.hooks_context
+    selection:
+        inherited_members: true
diff --git a/docs/api/initializers/Constant.md b/docs/api/initializers/Constant.md
@@ -1,5 +1,6 @@
 
 # elegy.initializers.Constant
+
 ::: elegy.initializers.Constant
     selection:
         inherited_members: true
diff --git a/docs/api/initializers/Orthogonal.md b/docs/api/initializers/Orthogonal.md
@@ -1,5 +1,6 @@
 
 # elegy.initializers.Orthogonal
+
 ::: elegy.initializers.Orthogonal
     selection:
         inherited_members: true
diff --git a/docs/api/initializers/RandomNormal.md b/docs/api/initializers/RandomNormal.md
@@ -1,5 +1,6 @@
 
 # elegy.initializers.RandomNormal
+
 ::: elegy.initializers.RandomNormal
     selection:
         inherited_members: true
diff --git a/docs/api/initializers/RandomUniform.md b/docs/api/initializers/RandomUniform.md
@@ -1,5 +1,6 @@
 
 # elegy.initializers.RandomUniform
+
 ::: elegy.initializers.RandomUniform
     selection:
         inherited_members: true
diff --git a/docs/api/initializers/TruncatedNormal.md b/docs/api/initializers/TruncatedNormal.md
@@ -1,5 +1,6 @@
 
 # elegy.initializers.TruncatedNormal
+
 ::: elegy.initializers.TruncatedNormal
     selection:
         inherited_members: true
diff --git a/docs/api/initializers/UniformScaling.md b/docs/api/initializers/UniformScaling.md
@@ -1,5 +1,6 @@
 
 # elegy.initializers.UniformScaling
+
 ::: elegy.initializers.UniformScaling
     selection:
         inherited_members: true
diff --git a/docs/api/initializers/VarianceScaling.md b/docs/api/initializers/VarianceScaling.md
@@ -1,5 +1,6 @@
 
 # elegy.initializers.VarianceScaling
+
 ::: elegy.initializers.VarianceScaling
     selection:
         inherited_members: true
diff --git a/docs/api/is_training.md b/docs/api/is_training.md
@@ -0,0 +1,6 @@
+
+# elegy.is_training
+
+::: elegy.module.is_training
+    selection:
+        inherited_members: true