updated.

docs/index.md

@@ -2,10 +2,10 @@
 
 **Red Coast** (redco) is a lightweight and user-friendly tool designed to automate distributed training and inference for large models while simplifying the ML pipeline development process without necessitating MLSys expertise from users.
 
-RedCoast supports *Large Models* + *Complex Algorithms*, in a *lightweight* and *user-friendly* manner: 
+RedCoast supports *Large Models* + *Complex Algorithms*, in a *lightweight* and *user-friendly* way: 
 
 * Large Models beyond Transformers, e.g, [Stable Diffusion](https://github.com/tanyuqian/redco/tree/master/examples/text_to_image), etc.
-* Complex algorithms beyond cross entropy, e.g., [Meta Learning](https://github.com/tanyuqian/redco/tree/master/examples/meta_learning), etc.
+* Complex algorithms beyond cross entropy, e.g., [Meta Learning](https://github.com/tanyuqian/redco/tree/master/examples/meta_learning), [DP Training](https://github.com/tanyuqian/redco/tree/master/examples/differential_private_training), etc.
 
 With RedCoast, to define a ML pipeline, only three functions are needed:
 

docs/mnist.md

@@ -1,4 +1,4 @@
-This is a trivial MNIST example with RedCoast. Runnable by 
+This is a trivial MNIST example with RedCoast (`pip install redco==0.4.22`). Runnable by 
 ```
 python main.py
 ```
@@ -47,14 +47,14 @@ def collate_fn(examples):
 
 
 # Loss function converting model inputs to a scalar loss
-def loss_fn(train_rng, state, params, batch, is_training):
+def loss_fn(rng, state, params, batch, is_training):
     logits = state.apply_fn({'params': params}, batch['images'])
     return optax.softmax_cross_entropy_with_integer_labels(
         logits=logits, labels=batch['labels']).mean()
 
 
 # Predict function converting model inputs to the model outputs
-def pred_fn(pred_rng, params, batch, model):
+def pred_fn(rng, params, batch, model):
     accs = model.apply({'params': params}, batch['images']).argmax(axis=-1)
     return {'acc': accs}
 

redco/__init__.py

@@ -12,7 +12,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-__version__ = '0.4.21'
+__version__ = '0.4.22'
 
 from .deployers import *
 from .trainers import *

redco/deployers/deployer.py

@@ -276,6 +276,9 @@ def gen_rng(self):
         return new_rng
 
     def gen_model_step_rng(self):
+        """Get a new random number generator key for distributed model step and
+        update the random state.
+        """
         rng = self.gen_rng()
         if self.mesh is None:
             rng = jax.random.split(

setup.py

@@ -17,7 +17,7 @@
 
 setup(
     name="redco",
-    version="0.4.21",
+    version="0.4.22",
     author="Bowen Tan",
     packages=find_packages(),
     install_requires=['jax', 'flax', 'optax'],