doc update module, functional, geometry types

Author: chrischoy

.github/workflows/docs.yml

@@ -29,16 +29,13 @@ jobs:
 
       - name: Install dependencies
         run: |
-          pip install mkdocs
-          pip install mkdocs-rtd-dropdown
-          pip install mkdocstrings
-          pip install mkdocstrings-python
-          pip install pymdown-extensions
-          pip install torch
-          pip install warpconvnet==0.3.4
+          python -m pip install --upgrade pip
+          pip install -r docs/requirements.txt
+          pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
+          echo "PYTHONPATH=$PWD" >> $GITHUB_ENV
 
       - name: Build documentation
-        run: mkdocs build -f docs/mkdocs-readthedocs.yml
+        run: mkdocs build
 
       - name: Setup Pages
         uses: actions/configure-pages@v4
@@ -58,4 +55,4 @@ jobs:
     steps:
       - name: Deploy to GitHub Pages
         id: deployment
-        uses: actions/deploy-pages@v4 
+        uses: actions/deploy-pages@v4

README.md

@@ -180,10 +180,10 @@ pytest tests/ --benchmark-only
 uv pip install -r docs/requirements.txt
 
 # Build docs
-mkdocs build -f mkdocs-readthedocs.yml
+mkdocs build
 
 # Serve locally
-mkdocs serve -f mkdocs-readthedocs.yml
+mkdocs serve
 ```
 
 📖 **Documentation**: [https://nvlabs.github.io/WarpConvNet/](https://nvlabs.github.io/WarpConvNet/)

docs/api/geometry.md

@@ -1,3 +1,56 @@
 # Geometry
 
 ::: warpconvnet.geometry
+
+## Geometry containers
+
+Every geometry type in WarpConvNet wraps a `Coords` instance and a `Features`
+instance that share the same ragged-batch metadata.
+
+- `warpconvnet.geometry.base.coords.Coords` stores concatenated coordinates plus an
+  `offsets` vector marking where each example begins.
+- `warpconvnet.geometry.base.features.Features` (and the `CatFeatures` and `PadFeatures`
+  specializations) stores feature tensors that obey the same offsets so coordinates and
+  features always stay aligned.
+- `warpconvnet.geometry.base.geometry.Geometry` wires the pair together, validates their
+  shapes, and exposes device/dtype utilities with AMP-aware accessors.
+
+This shared contract lets subclasses freely switch between point clouds, voxels, or grids
+without duplicating batching logic. See [Batched coordinate layout](./geometry_batched.md)
+for a deeper explanation of how concatenated tensors and offsets interact.
+
+## Types
+
+WarpConvNet ships several geometry containers that unify coordinate systems
+with their associated features. Use these types as the canonical interfaces for
+points, voxels, dense grids, and FIGConvNet factor grids.
+
+### Points
+
+Flexible point-cloud geometry supporting ragged batches, feature paddings, and
+neighbor search utilities for sparse convolution modules.
+
+::: warpconvnet.geometry.types.points.Points
+
+### Voxels
+
+Sparse voxel geometry that accepts integer coordinates with tensor strides and
+offers helpers to move between dense tensors and CSR-style batched features.
+
+::: warpconvnet.geometry.types.voxels.Voxels
+
+### Grid
+
+Regular dense grid representation that keeps `GridCoords` and `GridFeatures`
+in sync, providing utilities for shape validation, format conversions, and
+batch-aware initialization.
+
+::: warpconvnet.geometry.types.grid.Grid
+
+### Factor grid
+
+Container that bundles multiple `Grid` instances with distinct factorized
+memory formats so that FIGConvNet layers can operate on complementary spatial
+perspectives.
+
+::: warpconvnet.geometry.types.factor_grid.FactorGrid

docs/api/geometry_batched.md

@@ -0,0 +1,61 @@
+# Batched coordinate layout
+
+::: warpconvnet.geometry.base.geometry.Geometry
+
+Geometry containers in WarpConvNet treat batched data as a pair of concatenated tensors
+plus an `offsets` vector. This “batched” layout stores every sample back-to-back inside a
+single tensor while recording where each sample begins so ragged batches stay addressable.
+
+## Concatenated coordinates
+
+`warpconvnet.geometry.base.coords.Coords` keeps two tensors:
+
+- `batched_tensor`: shape `[N, D]` with all coordinates concatenated in batch order.
+- `offsets`: shape `[B + 1]` with `offsets[b]` marking the starting row for batch `b`.
+
+The difference `offsets[b + 1] - offsets[b]` gives the number of coordinates in sample
+`b`. Because the data lives in one tensor, it is cheap to move, sort, or feed to CUDA
+kernels that expect CSR-style layouts.
+
+## Features that share offsets
+
+`warpconvnet.geometry.base.features.Features` mirrors the coordinate batching strategy.
+Both `CatFeatures` (concatenated) and `PadFeatures` (padded) expose the same offsets so
+`Geometry` subclasses can hop between dense and ragged views without recomputing metadata.
+
+- Concatenated features: `[N, C]` storage with the same CSR-style offsets as the
+  coordinates.
+- Padded features: `[B, L_max, C]` storage when kernels need dense memory, but each row
+  still references the shared offsets to keep track of valid entries.
+
+## Putting it together
+
+The base `Geometry` class validates that coordinates and features have identical offsets,
+then exposes helpers such as `geometry.batch_indexed_coordinates` and AMP-aware feature
+accessors. Subclasses (e.g., `Points`, `Voxels`, `Grid`) inherit these guarantees and add
+type-specific methods on top.
+
+```python
+import torch
+from warpconvnet.geometry.types.points import Points
+
+coords = torch.tensor(
+    [
+        [0, 0, 0],
+        [0, 1, 0],
+        [1, 0, 0],
+        [5, 5, 5],
+        [6, 5, 5],
+    ],
+    dtype=torch.int32,
+)
+offsets = torch.tensor([0, 3, 5], dtype=torch.int32)
+features = torch.randn(5, 32)
+
+points = Points(coords, features, offsets=offsets)
+# points.batched_coordinates.offsets == points.batched_features.offsets
+```
+
+In this example two ragged samples (lengths 3 and 2) share one coordinate tensor and one
+feature tensor. The shared offsets allow the geometry module to slice, move devices, or
+convert between concatenated and padded features without copying metadata.

docs/api/nn.md

@@ -1,3 +1,65 @@
 # Neural Networks
 
 ::: warpconvnet.nn
+
+## Modules
+
+### Activations
+
+::: warpconvnet.nn.modules.activations
+
+### Attention
+
+::: warpconvnet.nn.modules.attention
+
+### Base module
+
+::: warpconvnet.nn.modules.base_module
+
+### Factor grid
+
+::: warpconvnet.nn.modules.factor_grid
+
+### Grid convolution
+
+::: warpconvnet.nn.modules.grid_conv
+
+### MLP
+
+::: warpconvnet.nn.modules.mlp
+
+### Normalizations
+
+::: warpconvnet.nn.modules.normalizations
+
+### Point convolution
+
+::: warpconvnet.nn.modules.point_conv
+
+### Point pooling
+
+::: warpconvnet.nn.modules.point_pool
+
+### Prune
+
+::: warpconvnet.nn.modules.prune
+
+### Sequential
+
+::: warpconvnet.nn.modules.sequential
+
+### Sparse convolution
+
+::: warpconvnet.nn.modules.sparse_conv
+
+### Sparse depthwise convolution
+
+::: warpconvnet.nn.modules.sparse_conv_depth
+
+### Sparse pooling
+
+::: warpconvnet.nn.modules.sparse_pool
+
+### Transforms
+
+::: warpconvnet.nn.modules.transforms

docs/api/nn_functional.md

@@ -0,0 +1,61 @@
+# Neural Network Functionals
+
+::: warpconvnet.nn.functional
+
+## Functionals
+
+### Sparse convolution
+
+::: warpconvnet.nn.functional.sparse_conv
+
+### Sparse depthwise convolution
+
+::: warpconvnet.nn.functional.sparse_conv_depth
+
+### Sparse pooling
+
+::: warpconvnet.nn.functional.sparse_pool
+
+### Sparse ops helpers
+
+::: warpconvnet.nn.functional.sparse_ops
+
+### Point pooling
+
+::: warpconvnet.nn.functional.point_pool
+
+### Point unpooling
+
+::: warpconvnet.nn.functional.point_unpool
+
+### Grid convolution
+
+::: warpconvnet.nn.functional.grid_conv
+
+### Factor grid
+
+::: warpconvnet.nn.functional.factor_grid
+
+### Global pooling
+
+::: warpconvnet.nn.functional.global_pool
+
+### Feature transforms
+
+::: warpconvnet.nn.functional.transforms
+
+### Encodings
+
+::: warpconvnet.nn.functional.encodings
+
+### Normalizations
+
+::: warpconvnet.nn.functional.normalizations
+
+### Segmented arithmetic
+
+::: warpconvnet.nn.functional.segmented_arithmetics
+
+### Batched matrix multiplication
+
+::: warpconvnet.nn.functional.bmm

docs/deployment.md

@@ -9,17 +9,20 @@ The documentation is automatically built and deployed using GitHub Actions whene
 ### How it Works
 
 1. **GitHub Actions Workflow**: `.github/workflows/docs.yml`
+
    - Triggers on pushes to main/master branch
    - Installs Python dependencies
    - Builds documentation using MkDocs
    - Deploys to GitHub Pages
 
-2. **Configuration**: `mkdocs-readthedocs.yml`
+2. **Configuration**: `mkdocs.yml`
+
    - Uses ReadTheDocs theme
    - Configures site metadata
    - Sets up navigation structure
 
 3. **Dependencies**: `docs/requirements.txt`
+
    - Lists all required Python packages
    - Ensures consistent builds
 
@@ -32,7 +35,7 @@ If you need to deploy manually:
 uv pip install -r docs/requirements.txt
 
 # Build documentation
-mkdocs build -f mkdocs-readthedocs.yml
+mkdocs build
 
 # The site/ directory contains the built documentation
 ```
@@ -43,7 +46,7 @@ For local development and testing:
 
 ```bash
 # Serve documentation locally
-mkdocs serve -f mkdocs-readthedocs.yml
+mkdocs serve
 
 # This will start a local server at http://127.0.0.1:8000
 ```
@@ -60,16 +63,19 @@ To enable GitHub Pages:
 ## Troubleshooting
 
 ### Build Failures
+
 - Check GitHub Actions logs for error details
 - Ensure all dependencies are listed in `docs/requirements.txt`
 - Verify MkDocs configuration syntax
 
 ### Missing Pages
+
 - Check that all referenced markdown files exist
-- Verify navigation structure in `mkdocs-readthedocs.yml`
+- Verify navigation structure in `mkdocs.yml`
 - Ensure files are committed to the repository
 
 ### Theme Issues
+
 - ReadTheDocs theme is automatically installed by the workflow
-- Check theme configuration in `mkdocs-readthedocs.yml`
-- Verify JavaScript and CSS assets are loading correctly 
+- Check theme configuration in `mkdocs.yml`
+- Verify JavaScript and CSS assets are loading correctly