new: add the first structure of the project (#1)

Author: monatis

.devcontainer/Dockerfile

@@ -0,0 +1,50 @@
+FROM nvidia/cuda:12.4.0-runtime-ubuntu22.04 
+
+# Avoid prompts from apt
+ENV DEBIAN_FRONTEND=noninteractive
+
+# Install Python and other dependencies
+RUN apt-get update && apt-get install -y \
+    python3.10 \
+    python3.10-dev \
+    python3-pip \
+    python3-venv \
+    git \
+    curl \
+    wget \
+    build-essential \
+    && rm -rf /var/lib/apt/lists/*
+
+# Create symbolic links for python
+RUN ln -sf /usr/bin/python3.10 /usr/bin/python && \
+    ln -sf /usr/bin/python3.10 /usr/bin/python3
+
+# Create a non-root user
+ARG USERNAME=vscode
+ARG USER_UID=1000
+ARG USER_GID=$USER_UID
+
+RUN groupadd --gid $USER_GID $USERNAME \
+    && useradd --uid $USER_UID --gid $USER_GID -m $USERNAME \
+    && apt-get update \
+    && apt-get install -y sudo \
+    && echo $USERNAME ALL=\(root\) NOPASSWD:ALL > /etc/sudoers.d/$USERNAME \
+    && chmod 0440 /etc/sudoers.d/$USERNAME
+
+# Set up Python environment
+RUN python -m pip install --upgrade pip setuptools wheel
+
+# Install JAX with CUDA support
+RUN pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+
+# Install matplotlib for benchmarking script
+RUN pip install matplotlib
+
+# Set working directory
+WORKDIR /workspace
+
+# Switch to non-root user
+USER $USERNAME
+
+# Note: The project dependencies will be installed by the postCreateCommand in devcontainer.json
+# which runs `pip install -e '.[dev]'` after the container is created 

.devcontainer/README.md

@@ -0,0 +1,50 @@
+# Development Container for JAX Layers
+
+This directory contains configuration files for setting up a development container with JAX and CUDA support, which is especially useful for Windows users where JAX doesn't natively support CUDA.
+
+## Prerequisites
+
+To use this development container, you need:
+
+1. [Docker Desktop](https://www.docker.com/products/docker-desktop/) installed and configured with WSL 2 backend
+2. [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html) installed
+3. [Visual Studio Code](https://code.visualstudio.com/) with the [Remote - Containers](https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-containers) extension
+
+## GPU Support
+
+The container is configured to use all available GPUs. Make sure your NVIDIA drivers are up-to-date and that Docker has access to your GPUs.
+
+## Usage
+
+1. Open the project in Visual Studio Code
+2. Click on the green icon in the bottom-left corner of VS Code
+3. Select "Reopen in Container" from the menu
+4. Wait for the container to build and start (this may take a while the first time)
+
+Once the container is running, you'll have a fully configured development environment with:
+
+- Python 3.10
+- CUDA 12.2 with cuDNN 9
+- JAX with CUDA support
+- All dependencies from pyproject.toml
+
+## Dependency Management
+
+The container installs dependencies directly from your project's `pyproject.toml` file using the `pip install -e '.[dev]'` command, ensuring consistency between your development environment and the container.
+
+## Customization
+
+You can customize the container by modifying:
+
+- `devcontainer.json`: VS Code settings, extensions, and container configuration
+- `Dockerfile`: Base image, dependencies, and environment setup
+
+## Troubleshooting
+
+If you encounter issues with GPU access:
+
+1. Verify that Docker Desktop is configured to use WSL 2
+2. Check that NVIDIA Container Toolkit is properly installed
+3. Ensure your NVIDIA drivers are up-to-date
+4. Run `nvidia-smi` in WSL to verify GPU access
+5. Check Docker logs for any error messages related to GPU access

.devcontainer/devcontainer.json

@@ -0,0 +1,36 @@
+{
+    "name": "JAX Layers Development",
+    "build": {
+        "dockerfile": "Dockerfile",
+        "context": ".."
+    },
+    "runArgs": [
+        "--gpus=all"
+    ],
+    "customizations": {
+        "vscode": {
+            "extensions": [
+                "ms-python.python",
+                "ms-python.vscode-pylance",
+                "charliermarsh.ruff",
+                "matangover.mypy"
+            ],
+            "settings": {
+                "python.defaultInterpreterPath": "/usr/local/bin/python",
+                "python.linting.enabled": true,
+                "editor.formatOnSave": true,
+                "editor.codeActionsOnSave": {
+                    "source.organizeImports": "true",
+                    "source.fixAll": "true"
+                },
+                "python.formatting.provider": "none",
+                "[python]": {
+                    "editor.defaultFormatter": "charliermarsh.ruff"
+                }
+            }
+        }
+    },
+    "forwardPorts": [],
+    "postCreateCommand": "pip install -e '.[dev]'",
+    "remoteUser": "vscode"
+}

.devcontainer/verify_jax_cuda.py

@@ -0,0 +1,104 @@
+#!/usr/bin/env python3
+"""
+Script to verify that JAX with CUDA is working correctly.
+Run this script after the container is built to confirm GPU access.
+"""
+
+import re
+import subprocess
+import time
+
+import jax
+import jax.numpy as jnp
+
+
+def get_cuda_version():
+    """Get the CUDA version from nvcc."""
+    try:
+        result = subprocess.run(["nvcc", "--version"], capture_output=True, text=True)
+        version_match = re.search(r"release (\d+\.\d+)", result.stdout)
+        if version_match:
+            return version_match.group(1)
+        return "Unknown"
+    except Exception:
+        return "Unknown"
+
+
+def main():
+    print("\n" + "=" * 50)
+    print("JAX CUDA Verification Script")
+    print("=" * 50)
+
+    # Check JAX version
+    print(f"JAX version: {jax.__version__}")
+
+    # Check CUDA version
+    cuda_version = get_cuda_version()
+    print(f"CUDA version: {cuda_version}")
+
+    # Check available devices
+    print("\nAvailable devices:")
+    for i, device in enumerate(jax.devices()):
+        print(f"  Device {i}: {device}")
+
+    # Check if GPU is available
+    gpu_available = any(d.platform == "gpu" for d in jax.devices())
+    print(f"\nGPU available: {gpu_available}")
+
+    if not gpu_available:
+        print("\n⚠️  No GPU devices found! JAX is not using CUDA.")
+        print("Please check your installation and GPU configuration.")
+        return
+
+    # Run a simple benchmark
+    print("\nRunning simple matrix multiplication benchmark...")
+
+    # Create large matrices
+    n = 5000
+    print(f"Creating {n}x{n} matrices...")
+
+    # CPU benchmark
+    with jax.devices("cpu")[0]:
+        x_cpu = jnp.ones((n, n))
+        y_cpu = jnp.ones((n, n))
+
+        # Warm-up
+        _ = jnp.dot(x_cpu, y_cpu)
+        jax.block_until_ready(_)
+
+        # Benchmark
+        start = time.time()
+        result_cpu = jnp.dot(x_cpu, y_cpu)
+        jax.block_until_ready(result_cpu)
+        cpu_time = time.time() - start
+
+    # GPU benchmark
+    with jax.devices("gpu")[0]:
+        x_gpu = jnp.ones((n, n))
+        y_gpu = jnp.ones((n, n))
+
+        # Warm-up
+        _ = jnp.dot(x_gpu, y_gpu)
+        jax.block_until_ready(_)
+
+        # Benchmark
+        start = time.time()
+        result_gpu = jnp.dot(x_gpu, y_gpu)
+        jax.block_until_ready(result_gpu)
+        gpu_time = time.time() - start
+
+    # Print results
+    print(f"\nCPU time: {cpu_time:.4f} seconds")
+    print(f"GPU time: {gpu_time:.4f} seconds")
+    print(f"Speedup: {cpu_time / gpu_time:.2f}x")
+
+    if cpu_time > gpu_time:
+        print("\n✅ GPU is faster than CPU! JAX with CUDA is working correctly.")
+    else:
+        print("\n⚠️  GPU is not faster than CPU. Something might be wrong with the CUDA setup.")
+
+    print("\n" + "=" * 50)
+
+
+if __name__ == "__main__":
+    main()

.github/workflows/docs.yml

@@ -0,0 +1,36 @@
+name: Documentation
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  docs:
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+    
+    - name: Set up Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: "3.10"
+        
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install -e ".[dev]"
+        
+    - name: Build documentation
+      run: |
+        cd docs
+        make html
+        
+    - name: Deploy to GitHub Pages
+      if: github.event_name == 'push' && github.ref == 'refs/heads/main'
+      uses: peaceiris/actions-gh-pages@v3
+      with:
+        github_token: ${{ secrets.GITHUB_TOKEN }}
+        publish_dir: ./docs/_build/html 

.github/workflows/mypy.yml

@@ -0,0 +1,29 @@
+name: Type Checking
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  mypy:
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+    
+    - name: Set up Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: "3.10"
+        
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install -e ".[dev]"
+        pip install mypy
+        
+    - name: Run mypy
+      run: |
+        mypy jax_layers tests 

.github/workflows/ruff.yml

@@ -0,0 +1,29 @@
+name: Ruff
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  ruff:
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+    
+    - name: Set up Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: "3.10"
+        
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install ruff
+        
+    - name: Run Ruff
+      run: |
+        ruff check .
+        ruff format --check . 

.github/workflows/tests.yml

@@ -0,0 +1,31 @@
+name: Tests
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.11", "3.12"]
+
+    steps:
+    - uses: actions/checkout@v4
+    
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python-version }}
+        
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install -e ".[dev]"
+        
+    - name: Run tests
+      run: |
+        python tests/run_tests.py 

.python-version

@@ -0,0 +1 @@
+3.12

docs/.gitignore

@@ -0,0 +1 @@
+_build/