Deployed fa89e21 with MkDocs version: 1.6.1

tanyuqian · Oct 20, 2024 · 2860cd1 · 2860cd1
1 parent 50413a3
commit 2860cd1
Show file tree

Hide file tree

Showing 5 changed files with 275 additions and 151 deletions.
diff --git a/docs/deployer/index.html b/docs/deployer/index.html
@@ -1869,7 +1869,7 @@ <h2 id="redco.deployers.deployer.Deployer" class="doc doc-heading">
 
 <span class="sd">        Args:</span>
 <span class="sd">            ckpt_dir (str): Directory of the checkpoint.</span>
-<span class="sd">            params_sharding_rules (list[tuple]): Sharding rules for parameters.</span>
+<span class="sd">            params_sharding_rules (PyTree): Sharding rules for parameters.</span>
 <span class="sd">            optimizer (optax optimizer): Optimizer for loading opt_state.</span>
 <span class="sd">            float_dtype (`jax.numpy.dtype`): Dtype for floating point numbers.</span>
 <span class="sd">            load_params (bool): Whether to load the parameters.</span>
@@ -3185,7 +3185,7 @@ <h3 id="redco.deployers.deployer.Deployer.load_ckpt" class="doc doc-heading">
           <tr class="doc-section-item">
             <td><code>params_sharding_rules</code></td>
             <td>
-                  <code>list[tuple]</code>
+                  <code>PyTree</code>
             </td>
             <td>
               <div class="doc-md-description">
@@ -3372,7 +3372,7 @@ <h3 id="redco.deployers.deployer.Deployer.load_ckpt" class="doc doc-heading">
 
 <span class="sd">    Args:</span>
 <span class="sd">        ckpt_dir (str): Directory of the checkpoint.</span>
-<span class="sd">        params_sharding_rules (list[tuple]): Sharding rules for parameters.</span>
+<span class="sd">        params_sharding_rules (PyTree): Sharding rules for parameters.</span>
 <span class="sd">        optimizer (optax optimizer): Optimizer for loading opt_state.</span>
 <span class="sd">        float_dtype (`jax.numpy.dtype`): Dtype for floating point numbers.</span>
 <span class="sd">        load_params (bool): Whether to load the parameters.</span>