fix up jax tutorial and bump version

AlexanderFengler · Dec 8, 2024 · 725630c · 725630c
1 parent d96eb8b
commit 725630c
Show file tree

Hide file tree

Showing 22 changed files with 609 additions and 1,237 deletions.
diff --git a/docs/basic_tutorial/basic_tutorial_jax.ipynb b/docs/basic_tutorial/basic_tutorial_jax.ipynb
diff --git a/docs/basic_tutorial/data/jax_models/angle/_lan_angle__data_details.pickle b/docs/basic_tutorial/data/jax_models/angle/_lan_angle__data_details.pickle
diff --git a/docs/basic_tutorial/data/jax_models/angle/_lan_angle__jax_training_history.csv b/docs/basic_tutorial/data/jax_models/angle/_lan_angle__jax_training_history.csv
@@ -0,0 +1,6 @@
+,epoch,val_loss
+0,0.0,0.29503539204597473
+1,1.0,0.21996137499809265
+2,2.0,0.1790909618139267
+3,3.0,0.10524507611989975
+4,4.0,0.07911499589681625
diff --git a/docs/basic_tutorial/data/jax_models/angle/_lan_angle__train_config.pickle b/docs/basic_tutorial/data/jax_models/angle/_lan_angle__train_config.pickle
diff --git a/docs/basic_tutorial/data/jax_models/angle/_lan_angle__train_state.jax b/docs/basic_tutorial/data/jax_models/angle/_lan_angle__train_state.jax
diff --git a/docs/basic_tutorial/data/jax_models/angle/jax_lan_angle__data_details.pickle b/docs/basic_tutorial/data/jax_models/angle/jax_lan_angle__data_details.pickle
diff --git a/docs/basic_tutorial/data/jax_models/angle/jax_lan_angle__jax_training_history.csv b/docs/basic_tutorial/data/jax_models/angle/jax_lan_angle__jax_training_history.csv
@@ -0,0 +1,6 @@
+,epoch,val_loss
+0,0.0,0.3224255442619324
+1,1.0,0.24499981105327606
+2,2.0,0.24103233218193054
+3,3.0,0.1361701339483261
+4,4.0,0.11248210817575455
diff --git a/docs/basic_tutorial/data/jax_models/angle/jax_lan_angle__train_config.pickle b/docs/basic_tutorial/data/jax_models/angle/jax_lan_angle__train_config.pickle
diff --git a/docs/basic_tutorial/data/jax_models/angle/jax_lan_angle__train_state.jax b/docs/basic_tutorial/data/jax_models/angle/jax_lan_angle__train_state.jax
diff --git a/docs/basic_tutorial/data/jax_models/angle/jax_network_config.pickle b/docs/basic_tutorial/data/jax_models/angle/jax_network_config.pickle
diff --git a/docs/basic_tutorial/data/jax_models/angle/lan_angle__train_state.jax b/docs/basic_tutorial/data/jax_models/angle/lan_angle__train_state.jax
diff --git a/docs/overrides/main.html b/docs/overrides/main.html
@@ -8,7 +8,7 @@
     Navigate the site here!
 </span>
 <span class="right-margin">
-    v0.4.7 is out!
+    v0.4.8 is out!
 </span>
 <span>
     <span class="twemoji">

diff --git a/lan_pipeline_env.yml b/lan_pipeline_env.yml
@@ -92,7 +92,7 @@ dependencies:
       - jupyterlab-server==2.25.0
       - jupyterlab-widgets==3.0.9
       - kiwisolver==1.4.5
-      - lanfactory==0.4.7
+      - lanfactory==0.4.8
       - markdown-it-py==3.0.0
       - markupsafe==2.1.2
       - matplotlib==3.8.0

diff --git a/lanfactory/__init__.py b/lanfactory/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "0.4.7"
+__version__ = "0.4.8"
 
 from . import config
 from . import trainers

diff --git a/lanfactory/trainers/jax_mlp.py b/lanfactory/trainers/jax_mlp.py
@@ -129,13 +129,10 @@ def __call__(self, inputs):
                     x = self.activation_funs[i](x)
 
         if (not self.train) and (self.train_output_type == "logprob"):
-            print("passing through identity")
             x = x  # just for pedagogy
         elif (not self.train) and (self.train_output_type == "logits"):
-            print("passing through transform")
             x = -jnp.log((1 + jnp.exp(-x)))
         elif not self.train:
-            print("passing through identity 2")
             x = x  # just for pedagogy
 
         return x

diff --git a/mkdocs.yml b/mkdocs.yml
@@ -8,6 +8,8 @@ nav:
       - Overview: index.md
   - Basic Tutorial:
       - Installation: basic_tutorial/basic_tutorial.ipynb
+  - Basic Tutorial / JAX:
+      - Installation: basic_tutorial_jax/basic_tutorial_jax.ipynb
   - API:
       - lanfactory: api/lanfactory.md
       - config: api/config.md