Add simple Flax low-level API Model example to README.md

README.md

@@ -131,6 +131,31 @@ model.fit(
     callbacks=[elegy.callbacks.TensorBoard("summaries")]
 )
 ```
+#### Using Jax Frameworks
+It is straightforward to integrate other functional JAX libraries with this 
+low-level API:
+
+```python
+class LinearClassifier(elegy.Model):
+    def test_step(
+        self, x, y_true, states: elegy.States, initializing: bool, rng: elegy.RNGSeq
+    ):
+        x = jnp.reshape(x, (x.shape[0], -1)) / 255
+        if initializing:
+            params = self.module.init({'params': rng.next()}, x)['params']
+        else:
+            params = states.net_params
+
+        logits = self.module.apply(
+            {'params': params}, x, rngs={'params': rng.next(), 'dropout': rng.next()}
+        )
+        labels = jax.nn.one_hot(y_true, 10)
+        loss = jnp.mean(-jnp.sum(labels * jax.nn.log_softmax(logits), axis=-1))
+        accuracy = jnp.mean(jnp.argmax(logits, axis=-1) == y_true)
+
+        logs = dict(accuracy=accuracy, loss=loss)
+        return loss, logs, states.update(rng=rng, net_params=params)
+```
 
 ## More Info
 * [Getting Started: High-level API](https://poets-ai.github.io/elegy/getting-started-high-level-api/) tutorial.
@@ -165,4 +190,4 @@ year = {2020},
 ```
 
 
-Where the current *version* may be retrieved either from the `Release` tag or the file [elegy/\_\_init\_\_.py](https://github.com/poets-ai/elegy/blob/master/elegy/__init__.py) and the *year* corresponds to the project's release year.
+Where the current *version* may be retrieved either from the `Release` tag or the file [elegy/\_\_init\_\_.py](https://github.com/poets-ai/elegy/blob/master/elegy/__init__.py) and the *year* corresponds to the project's release year.