Updated README.

Author: vbharadwaj-bk

README.md

@@ -3,6 +3,7 @@
 [![License](https://img.shields.io/badge/License-BSD_3--Clause-blue.svg)](https://opensource.org/licenses/BSD-3-Clause)
 
 [[Examples]](#show-me-some-examples) 
+[[JAX Examples]](#jax-examples)
 [[Citation and Acknowledgements]](#citation-and-acknowledgements)
 
 OpenEquivariance is a CUDA and HIP kernel generator for the Clebsch-Gordon tensor product, 
@@ -29,48 +30,15 @@ computation and memory consumption significantly.
 For detailed instructions on tests, benchmarks, MACE / Nequip, and our API,
 check out the [documentation](https://passionlab.github.io/OpenEquivariance).
 
-⭐️ **JAX Support**: Our latest update brings
-support for JAX. To install, execute the following commands in order:
+⭐️ **JAX**: Our latest update brings
+support for JAX. To install, execute the following 
+commands in order:
 
 ```
 pip install openequivariance[jax]
 pip install openequivariance_extjax --no-build-isolation
 ```
-
-```python
-os.environ["OEQ_NOTORCH"] = "1"
-import openequivariance as oeq
-import jax
-
-seed = 42
-key = jax.random.PRNGKey(seed)
-
-node_ct, nonzero_ct = 3, 4
-edge_index = jax.numpy.array(
-    [
-        [0, 1, 1, 2],
-        [1, 0, 2, 1],
-    ],
-    dtype=jax.numpy.int32, # NOTE: This int32, not int64
-)
-
-X = jax.random.uniform(key, shape=(node_ct, X_ir.dim), minval=0.0, maxval=1.0, dtype=jax.numpy.float32)
-Y = jax.random.uniform(key, shape=(nonzero_ct, Y_ir.dim),
-                        minval=0.0, maxval=1.0, dtype=jax.numpy.float32)
-W = jax.random.uniform(key, shape=(nonzero_ct, problem.weight_numel),
-                        minval=0.0, maxval=1.0, dtype=jax.numpy.float32)
-
-# Reuse problem from earlier
-tp_conv = oeq.jax.TensorProductConv(problem, deterministic=False)
-Z = tp_conv.forward(
-    X, Y, W, edge_index[0], edge_index[1]
-)
-print(jax.numpy.linalg.norm(Z))
-```
-
-📣 📣 OpenEquivariance was accepted to the 2025 SIAM Conference on Applied and 
-Computational Discrete Algorithms (Proceedings Track)! Catch the talk in 
-Montréal and check out the [camera-ready copy on Arxiv](https://arxiv.org/abs/2501.13986) (available May 12, 2025).
+See below for example usage.
 
 ## Show me some examples
 Here's a CG tensor product implemented by e3nn: 
@@ -166,6 +134,45 @@ print(torch.norm(Z))
 `deterministic=False`, the `sender` and `receiver` indices can have
 arbitrary order.
 
+## JAX Examples
+After installation, use the library
+as follows. Set `OEQ_NOTORCH=1`
+in your environment to avoid the PyTorch import in
+the regular `openequivariance` package.
+```python
+import jax
+import os
+
+os.environ["OEQ_NOTORCH"] = "1"
+import openequivariance as oeq
+
+seed = 42
+key = jax.random.PRNGKey(seed)
+
+node_ct, nonzero_ct = 3, 4
+edge_index = jax.numpy.array(
+    [
+        [0, 1, 1, 2],
+        [1, 0, 2, 1],
+    ],
+    dtype=jax.numpy.int32, # NOTE: This int32, not int64
+)
+
+X = jax.random.uniform(key, shape=(node_ct, X_ir.dim), minval=0.0, maxval=1.0, dtype=jax.numpy.float32)
+Y = jax.random.uniform(key, shape=(nonzero_ct, Y_ir.dim),
+                        minval=0.0, maxval=1.0, dtype=jax.numpy.float32)
+W = jax.random.uniform(key, shape=(nonzero_ct, problem.weight_numel),
+                        minval=0.0, maxval=1.0, dtype=jax.numpy.float32)
+
+# Reuse problem from earlier
+# ...
+tp_conv = oeq.jax.TensorProductConv(problem, deterministic=False)
+Z = tp_conv.forward(
+    X, Y, W, edge_index[0], edge_index[1]
+)
+print(jax.numpy.linalg.norm(Z))
+```
+
 ## Citation and Acknowledgements
 If you find this code useful, please cite our paper: