Update readme, Add badge and cross links,

PiperOrigin-RevId: 466330294

Author: Conchylicultor

README.md

@@ -1,11 +1,22 @@
 # Dataclass Array
 
+[![Unittests](https://github.com/google-research/dataclass_array/actions/workflows/pytest_and_autopublish.yml/badge.svg)](https://github.com/google-research/visu3d/actions/workflows/pytest_and_autopublish.yml)
+[![PyPI version](https://badge.fury.io/py/dataclass_array.svg)](https://badge.fury.io/py/dataclass_array)
+
 `DataclassArray` are dataclasses which behave like numpy-like arrays (can be
-batched, reshaped, sliced,...), but are compatible with Jax, TensorFlow, and
-numpy (with torch support planned).
+batched, reshaped, sliced,...), compatible with Jax, TensorFlow, and numpy (with
+torch support planned).
+
+This reduce boilerplate and improve readability. See the
+[motivating examples](#motivating-examples) section bellow.
+
+To view an example of dataclass arrays used in practice, see
+[visu3d](https://github.com/google-research/visu3d).
 
 ## Documentation
 
+### Definition
+
 To create a `dca.DataclassArray`, take a frozen dataclass and:
 
 *   Inherit from `dca.DataclassArray`
@@ -23,6 +34,8 @@ class Ray(dca.DataclassArray):
   dir: FloatArray['*batch_shape 3']
 ```
 
+### Usage
+
 Afterwards, the dataclass can be used as a numpy array:
 
 ```python
@@ -75,6 +88,100 @@ class MyArray(dca.DataclassArray):
   f: np.array
 ```
 
+### Vectorization
+
+`@dca.vectorize_method` allow your dataclass method to automatically support
+batching:
+
+1.  Implement method as if `self.shape == ()`
+2.  Decorate the method with `dca.vectorize_method`
+
+```python
+@dataclasses.dataclass(frozen=True)
+class Camera(dca.DataclassArray):
+  K: FloatArray['*batch_shape 4 4']
+  resolution = tuple[int, int]
+
+  @dca.vectorize_method
+  def rays(self) -> Ray:
+    # Inside `@dca.vectorize_method` shape is always guarantee to be `()`
+    assert self.shape == ()
+    assert self.K.shape == (4, 4)
+
+    # Compute the ray as if there was only a single camera
+    return Ray(pos=..., dir=...)
+```
+
+Afterward, we can generate rays for multiple camera batched together:
+
+```python
+cams = Camera(K=K)  # K.shape == (num_cams, 4, 4)
+rays = cams.rays()  # Generate the rays for all the cameras
+
+cams.shape == (num_cams,)
+rays.shape == (num_cams, h, w)
+```
+
+`@dca.vectorize_method` is similar to `jax.vmap` but:
+
+*   Only work on `dca.DataclassArray` methods
+*   Instead of vectorizing a single axis, `@dca.vectorize_method` will vectorize
+    over `*self.shape` (not just `self.shape[0]`). This is like if `vmap` was
+    applied to `self.flatten()`
+*   When multiple arguments, axis with dimension `1` are brodcasted.
+
+For example, with `__matmul__(self, x: T) -> T`:
+
+```python
+() @ (*x,) -> (*x,)
+(b,) @ (b, *x) -> (b, *x)
+(b,) @ (1, *x) -> (b, *x)
+(1,) @ (b, *x) -> (b, *x)
+(b, h, w) @ (b, h, w, *x) -> (b, h, w, *x)
+(1, h, w) @ (b, 1, 1, *x) -> (b, h, w, *x)
+(a, *x) @ (b, *x) -> Error: Incompatible a != b
+```
+
+To test on Colab, see the `visu3d`
+dataclass [Colab tutorial](https://colab.research.google.com/github/google-research/visu3d/blob/main/docs/dataclass.ipynb).
+
+## Motivating examples
+
+`dca.DataclassArray` improve readability by simplifying common patterns:
+
+*   Reshaping all fields of a dataclass:
+
+    Before (`rays` is simple `dataclass`):
+
+    ```python
+    num_rays = math.prod(rays.origins.shape[:-1])
+    rays = jax.tree_map(lambda r: r.reshape((num_rays, -1)), rays)
+    ```
+
+    After (`rays` is `DataclassArray`):
+
+    ```python
+    rays = rays.flatten()  # (b, h, w) -> (b*h*w,)
+    ```
+
+*   Rendering a video:
+
+    Before (`cams: list[Camera]`):
+
+    ```python
+    img = cams[0].render(scene)
+    imgs = np.stack([cam.render(scene) for cam in cams[::2]])
+    imgs = np.stack([cam.render(scene) for cam in cams])
+    ```
+
+    After (`cams: Camera` with `cams.shape == (num_cams,)`):
+
+    ```python
+    img = cams[0].render(scene)  # Render only the first camera (to debug)
+    imgs = cams[::2].render(scene)  # Render 1/2 frames (for quicker iteration)
+    imgs = cams.render(scene)  # Render all cameras at once
+    ```
+
 ## Installation
 
 ```sh