Ekfac

.github/workflows/build.yml

@@ -0,0 +1,39 @@
+name: build
+
+on:
+  push:
+    branches:
+      - ekfac
+  pull_request:
+    branches:
+      - ekfac
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e ".[dev,faiss]"
+      # TODO: Proper test infrastructure for tests/ekfac_tests
+      # - name: Run tests
+      #   run: pytest
+      # TODO: run pyright on whole codebase
+      - name: Type Checking bergson/hessians
+        uses: jakebailey/pyright-action@v1
+        with:
+          version: 1.1.406
+          working-directory: bergson/hessians
+      - name: Type Checking tests/ekfac_tests
+        uses: jakebailey/pyright-action@v1
+        with:
+          version: 1.1.406
+          working-directory: tests/ekfac_tests
+      - name: build
+        run: pip wheel --no-deps -w dist .
+env:
+  HF_HUB_DOWNLOAD_TIMEOUT: 100

.gitignore

@@ -161,23 +161,47 @@ dmypy.json
 cython_debug/
 
 # PyCharm
-#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
-#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
-#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  JetBrains specific template is maintained in a separate JetBrains. that can
+#  be found at https://github.com/github//blob/main/Global/JetBrains.
+#  and can be added to the global  or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
 
+# VS Code
+.vscode/
+
 # Ruff stuff:
 .ruff_cache/
 
 # PyPI configuration file
 .pypirc
 
+# models
+*.pt
+*.pth
+*.safetensors
+*.json
+*.jsonl
+*.txt
+*.arrow
+*.bin
+*.csv
+*.npy
+
+# plots
+*.png
+*.jpg
+*.jpeg
+*.gif
+
+# debugging results
+*.svg
+*.pickle
+# Faiss index files
+*.faiss
 # Local directory for run artifacts
 runs/
 cache/
 
 wandb/
 .vscode/
-
-

bergson/__main__.py

@@ -1,11 +1,16 @@
 from simple_parsing import parse
 
-from .build import build_gradient_dataset
+from bergson.distributed import distributed_computing
+
+from .collection import collect_gradients
 from .data import IndexConfig
 
 
 def main():
-    build_gradient_dataset(parse(IndexConfig))
+    distributed_computing(
+        parse(IndexConfig),
+        worker_fn=collect_gradients,
+    )
 
 
 if __name__ == "__main__":