update docs

Author: pdcook

.buildinfo

@@ -1,4 +1,4 @@
 # Sphinx build info version 1
 # This file records the configuration used when building these files. When it is not found, a full rebuild will be done.
-config: ec0cce5fd9708db0ba620f6383be6121
+config: 5d9106275669d325fa0100d5fad3379f
 tags: 645f666f9bcd5a90fca523b33c5a78b7

_modules/index.html

@@ -131,8 +131,8 @@
   <ul class="bd-navbar-elements navbar-nav">
 
 <li class="nav-item ">
-  <a class="nav-link nav-internal" href="../model.html">
-    Model
+  <a class="nav-link nav-internal" href="../models.html">
+    Models
   </a>
 </li>
 
@@ -231,8 +231,8 @@
   <ul class="bd-navbar-elements navbar-nav">
 
 <li class="nav-item ">
-  <a class="nav-link nav-internal" href="../model.html">
-    Model
+  <a class="nav-link nav-internal" href="../models.html">
+    Models
   </a>
 </li>
 
@@ -342,23 +342,24 @@ <h1>All modules for which code is available</h1>
 <li><a href="parametricmatrixmodels/modules/basemodule.html">parametricmatrixmodels.modules.basemodule</a></li>
 <li><a href="parametricmatrixmodels/modules/bias.html">parametricmatrixmodels.modules.bias</a></li>
 <li><a href="parametricmatrixmodels/modules/comment.html">parametricmatrixmodels.modules.comment</a></li>
+<li><a href="parametricmatrixmodels/modules/concatenateleaves.html">parametricmatrixmodels.modules.concatenateleaves</a></li>
+<li><a href="parametricmatrixmodels/modules/constant.html">parametricmatrixmodels.modules.constant</a></li>
 <li><a href="parametricmatrixmodels/modules/eigensystem.html">parametricmatrixmodels.modules.eigensystem</a></li>
 <li><a href="parametricmatrixmodels/modules/eigenvalues.html">parametricmatrixmodels.modules.eigenvalues</a></li>
 <li><a href="parametricmatrixmodels/modules/eigenvectors.html">parametricmatrixmodels.modules.eigenvectors</a></li>
+<li><a href="parametricmatrixmodels/modules/einsum.html">parametricmatrixmodels.modules.einsum</a></li>
+<li><a href="parametricmatrixmodels/modules/expectationvaluesum.html">parametricmatrixmodels.modules.expectationvaluesum</a></li>
 <li><a href="parametricmatrixmodels/modules/flatten.html">parametricmatrixmodels.modules.flatten</a></li>
-<li><a href="parametricmatrixmodels/modules/func.html">parametricmatrixmodels.modules.func</a></li>
-<li><a href="parametricmatrixmodels/modules/legacyaffineobservablepmm.html">parametricmatrixmodels.modules.legacyaffineobservablepmm</a></li>
+<li><a href="parametricmatrixmodels/modules/funcbase.html">parametricmatrixmodels.modules.funcbase</a></li>
 <li><a href="parametricmatrixmodels/modules/linearnn.html">parametricmatrixmodels.modules.linearnn</a></li>
-<li><a href="parametricmatrixmodels/modules/lowrankaffineeigenvaluepmm.html">parametricmatrixmodels.modules.lowrankaffineeigenvaluepmm</a></li>
-<li><a href="parametricmatrixmodels/modules/lowrankaffinehermitianmatrix.html">parametricmatrixmodels.modules.lowrankaffinehermitianmatrix</a></li>
-<li><a href="parametricmatrixmodels/modules/lowrankaffineobservablepmm.html">parametricmatrixmodels.modules.lowrankaffineobservablepmm</a></li>
-<li><a href="parametricmatrixmodels/modules/lowranktransitionamplitudesum.html">parametricmatrixmodels.modules.lowranktransitionamplitudesum</a></li>
-<li><a href="parametricmatrixmodels/modules/multimodule.html">parametricmatrixmodels.modules.multimodule</a></li>
-<li><a href="parametricmatrixmodels/modules/nonnegativelinearnn.html">parametricmatrixmodels.modules.nonnegativelinearnn</a></li>
+<li><a href="parametricmatrixmodels/modules/matmul.html">parametricmatrixmodels.modules.matmul</a></li>
 <li><a href="parametricmatrixmodels/modules/prelu.html">parametricmatrixmodels.modules.prelu</a></li>
 <li><a href="parametricmatrixmodels/modules/reshape.html">parametricmatrixmodels.modules.reshape</a></li>
-<li><a href="parametricmatrixmodels/modules/subsetmodule.html">parametricmatrixmodels.modules.subsetmodule</a></li>
 <li><a href="parametricmatrixmodels/modules/transitionamplitudesum.html">parametricmatrixmodels.modules.transitionamplitudesum</a></li>
+<li><a href="parametricmatrixmodels/modules/treeflatten.html">parametricmatrixmodels.modules.treeflatten</a></li>
+<li><a href="parametricmatrixmodels/modules/treekey.html">parametricmatrixmodels.modules.treekey</a></li>
+<li><a href="parametricmatrixmodels/nonsequentialmodel.html">parametricmatrixmodels.nonsequentialmodel</a></li>
+<li><a href="parametricmatrixmodels/sequentialmodel.html">parametricmatrixmodels.sequentialmodel</a></li>
 </ul>
 
                 </article>

_modules/parametricmatrixmodels/model.html

@@ -131,8 +131,8 @@
   <ul class="bd-navbar-elements navbar-nav">
 
 <li class="nav-item ">
-  <a class="nav-link nav-internal" href="../../../model.html">
-    Model
+  <a class="nav-link nav-internal" href="../../../models.html">
+    Models
   </a>
 </li>
 
@@ -231,8 +231,8 @@
   <ul class="bd-navbar-elements navbar-nav">
 
 <li class="nav-item ">
-  <a class="nav-link nav-internal" href="../../../model.html">
-    Model
+  <a class="nav-link nav-internal" href="../../../models.html">
+    Models
   </a>
 </li>
 
@@ -337,11 +337,14 @@
                 <article class="bd-article">
 
   <h1>Source code for parametricmatrixmodels.modules.comment</h1><div class="highlight"><pre>
-<span></span><span class="kn">from</span><span class="w"> </span><span class="nn">__future__</span><span class="w"> </span><span class="kn">import</span> <span class="n">annotations</span>
-
-<span class="kn">from</span><span class="w"> </span><span class="nn">typing</span><span class="w"> </span><span class="kn">import</span> <span class="n">Any</span><span class="p">,</span> <span class="n">Callable</span>
-
-<span class="kn">import</span><span class="w"> </span><span class="nn">jax.numpy</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">np</span>
+<span></span><span class="kn">from</span><span class="w"> </span><span class="nn">parametricmatrixmodels.typing</span><span class="w"> </span><span class="kn">import</span> <span class="p">(</span>
+    <span class="n">Any</span><span class="p">,</span>
+    <span class="n">DataShape</span><span class="p">,</span>
+    <span class="n">HyperParams</span><span class="p">,</span>
+    <span class="n">ModuleCallable</span><span class="p">,</span>
+    <span class="n">Params</span><span class="p">,</span>
+    <span class="n">State</span><span class="p">,</span>
+<span class="p">)</span>
 
 <span class="kn">from</span><span class="w"> </span><span class="nn">.basemodule</span><span class="w"> </span><span class="kn">import</span> <span class="n">BaseModule</span>
 
@@ -369,11 +372,9 @@ <h1>Source code for parametricmatrixmodels.modules.comment</h1><div class="highl
         <span class="bp">self</span><span class="o">.</span><span class="n">comment</span> <span class="o">=</span> <span class="n">comment</span></div>
 
 
-<div class="viewcode-block" id="Comment.name">
-<a class="viewcode-back" href="../../../generated/parametricmatrixmodels.modules.Comment.html#parametricmatrixmodels.modules.Comment.name">[docs]</a>
+    <span class="nd">@property</span>
     <span class="k">def</span><span class="w"> </span><span class="nf">name</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">str</span><span class="p">:</span>
-        <span class="k">return</span> <span class="sa">f</span><span class="s2">&quot;# </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">comment</span><span class="si">}</span><span class="s2">&quot;</span> <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">comment</span> <span class="k">else</span> <span class="s2">&quot;#&quot;</span></div>
-
+        <span class="k">return</span> <span class="sa">f</span><span class="s2">&quot;# </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">comment</span><span class="si">}</span><span class="s2">&quot;</span> <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">comment</span> <span class="k">else</span> <span class="s2">&quot;#&quot;</span>
 
 <div class="viewcode-block" id="Comment.is_ready">
 <a class="viewcode-back" href="../../../generated/parametricmatrixmodels.modules.Comment.html#parametricmatrixmodels.modules.Comment.is_ready">[docs]</a>
@@ -391,63 +392,46 @@ <h1>Source code for parametricmatrixmodels.modules.comment</h1><div class="highl
 <a class="viewcode-back" href="../../../generated/parametricmatrixmodels.modules.Comment.html#parametricmatrixmodels.modules.Comment._get_callable">[docs]</a>
     <span class="k">def</span><span class="w"> </span><span class="nf">_get_callable</span><span class="p">(</span>
         <span class="bp">self</span><span class="p">,</span>
-    <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Callable</span><span class="p">[</span>
-        <span class="p">[</span>
-            <span class="nb">tuple</span><span class="p">[</span><span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">,</span> <span class="o">...</span><span class="p">],</span>
-            <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">,</span>
-            <span class="nb">bool</span><span class="p">,</span>
-            <span class="nb">tuple</span><span class="p">[</span><span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">,</span> <span class="o">...</span><span class="p">],</span>
-            <span class="n">Any</span><span class="p">,</span>
-        <span class="p">],</span>
-        <span class="nb">tuple</span><span class="p">[</span><span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">[</span><span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">,</span> <span class="o">...</span><span class="p">]],</span>
-    <span class="p">]:</span>
-        <span class="k">return</span> <span class="k">lambda</span> <span class="n">params</span><span class="p">,</span> <span class="n">input_NF</span><span class="p">,</span> <span class="n">training</span><span class="p">,</span> <span class="n">state</span><span class="p">,</span> <span class="n">rng</span><span class="p">:</span> <span class="p">(</span>
-            <span class="n">input_NF</span><span class="p">,</span>  <span class="c1"># output is the same as input</span>
+    <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">ModuleCallable</span><span class="p">:</span>
+        <span class="k">return</span> <span class="k">lambda</span> <span class="n">params</span><span class="p">,</span> <span class="n">data</span><span class="p">,</span> <span class="n">training</span><span class="p">,</span> <span class="n">state</span><span class="p">,</span> <span class="n">rng</span><span class="p">:</span> <span class="p">(</span>
+            <span class="n">data</span><span class="p">,</span>  <span class="c1"># output is the same as input</span>
             <span class="n">state</span><span class="p">,</span>  <span class="c1"># state is unchanged</span>
         <span class="p">)</span></div>
 
 
 <div class="viewcode-block" id="Comment.compile">
 <a class="viewcode-back" href="../../../generated/parametricmatrixmodels.modules.Comment.html#parametricmatrixmodels.modules.Comment.compile">[docs]</a>
-    <span class="k">def</span><span class="w"> </span><span class="nf">compile</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">rng</span><span class="p">:</span> <span class="n">Any</span><span class="p">,</span> <span class="n">input_shape</span><span class="p">:</span> <span class="nb">tuple</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="o">...</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
+    <span class="k">def</span><span class="w"> </span><span class="nf">compile</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">rng</span><span class="p">:</span> <span class="n">Any</span><span class="p">,</span> <span class="n">input_shape</span><span class="p">:</span> <span class="n">DataShape</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
         <span class="k">pass</span></div>
 
 
 <div class="viewcode-block" id="Comment.get_output_shape">
 <a class="viewcode-back" href="../../../generated/parametricmatrixmodels.modules.Comment.html#parametricmatrixmodels.modules.Comment.get_output_shape">[docs]</a>
-    <span class="k">def</span><span class="w"> </span><span class="nf">get_output_shape</span><span class="p">(</span>
-        <span class="bp">self</span><span class="p">,</span> <span class="n">input_shape</span><span class="p">:</span> <span class="nb">tuple</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="o">...</span><span class="p">]</span>
-    <span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">tuple</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="o">...</span><span class="p">]:</span>
+    <span class="k">def</span><span class="w"> </span><span class="nf">get_output_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">input_shape</span><span class="p">:</span> <span class="n">DataShape</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">DataShape</span><span class="p">:</span>
         <span class="k">return</span> <span class="n">input_shape</span>  <span class="c1"># output shape is the same as input shape</span></div>
 
 
 <div class="viewcode-block" id="Comment.get_hyperparameters">
 <a class="viewcode-back" href="../../../generated/parametricmatrixmodels.modules.Comment.html#parametricmatrixmodels.modules.Comment.get_hyperparameters">[docs]</a>
-    <span class="k">def</span><span class="w"> </span><span class="nf">get_hyperparameters</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">]:</span>
+    <span class="k">def</span><span class="w"> </span><span class="nf">get_hyperparameters</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">HyperParams</span><span class="p">:</span>
         <span class="k">return</span> <span class="p">{</span><span class="s2">&quot;comment&quot;</span><span class="p">:</span> <span class="bp">self</span><span class="o">.</span><span class="n">comment</span><span class="p">}</span></div>
 
 
 <div class="viewcode-block" id="Comment.get_params">
 <a class="viewcode-back" href="../../../generated/parametricmatrixmodels.modules.Comment.html#parametricmatrixmodels.modules.Comment.get_params">[docs]</a>
-    <span class="k">def</span><span class="w"> </span><span class="nf">get_params</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">tuple</span><span class="p">[</span><span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">,</span> <span class="o">...</span><span class="p">]:</span>
+    <span class="k">def</span><span class="w"> </span><span class="nf">get_params</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Params</span><span class="p">:</span>
         <span class="k">return</span> <span class="p">()</span></div>
 
 
 <div class="viewcode-block" id="Comment.set_params">
 <a class="viewcode-back" href="../../../generated/parametricmatrixmodels.modules.Comment.html#parametricmatrixmodels.modules.Comment.set_params">[docs]</a>
-    <span class="k">def</span><span class="w"> </span><span class="nf">set_params</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">params</span><span class="p">:</span> <span class="nb">tuple</span><span class="p">[</span><span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">,</span> <span class="o">...</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
+    <span class="k">def</span><span class="w"> </span><span class="nf">set_params</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">params</span><span class="p">:</span> <span class="n">Params</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
         <span class="k">pass</span></div>
 
 
-<div class="viewcode-block" id="Comment.get_state">
-<a class="viewcode-back" href="../../../generated/parametricmatrixmodels.modules.Comment.html#parametricmatrixmodels.modules.Comment.get_state">[docs]</a>
-    <span class="k">def</span><span class="w"> </span><span class="nf">get_state</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">tuple</span><span class="p">[</span><span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">,</span> <span class="o">...</span><span class="p">]:</span>
-        <span class="k">return</span> <span class="p">()</span></div>
-
-
 <div class="viewcode-block" id="Comment.set_state">
 <a class="viewcode-back" href="../../../generated/parametricmatrixmodels.modules.Comment.html#parametricmatrixmodels.modules.Comment.set_state">[docs]</a>
-    <span class="k">def</span><span class="w"> </span><span class="nf">set_state</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">state</span><span class="p">:</span> <span class="nb">tuple</span><span class="p">[</span><span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">,</span> <span class="o">...</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
+    <span class="k">def</span><span class="w"> </span><span class="nf">set_state</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">state</span><span class="p">:</span> <span class="n">State</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
         <span class="k">pass</span></div>
 </div>