feat: add pre-commit hooks configuration with pyink and pylint

.pre-commit-config.yaml

@@ -0,0 +1,48 @@
+# Pre-commit hooks configuration
+# See https://pre-commit.com for more information
+# See https://pre-commit.com/hooks.html for more hooks
+
+default_language_version:
+  python: python3
+
+repos:
+  # Standard pre-commit hooks
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.5.0
+    hooks:
+      - id: trailing-whitespace
+        args: [--markdown-linebreak-ext=md]
+      - id: end-of-file-fixer
+      - id: check-yaml
+      - id: check-toml
+      - id: check-json
+      - id: check-added-large-files
+        args: ['--maxkb=1000']
+      - id: check-merge-conflict
+      - id: check-case-conflict
+      - id: mixed-line-ending
+        args: ['--fix=lf']
+      - id: debug-statements
+
+  # Python code formatting with pyink (Google's Python formatter)
+  - repo: https://github.com/google/pyink
+    rev: '23.10.0'
+    hooks:
+      - id: pyink
+        args:
+          - --line-length=80
+          - --pyink-indentation=2
+          - --pyink-use-majority-quotes
+
+  # Python linting with pylint
+  - repo: https://github.com/PyCQA/pylint
+    rev: v3.0.3
+    hooks:
+      - id: pylint
+        args:
+          - --rcfile=.pylintrc
+        additional_dependencies:
+          - jax
+          - jaxlib
+          - flax
+          - numpy

CONTRIBUTING.md

@@ -23,6 +23,68 @@ sign a new one.
 This project follows [Google's Open Source Community
 Guidelines](https://opensource.google/conduct/).
 
+## Development Environment Setup
+
+### Installing Pre-commit Hooks
+
+This repository uses [pre-commit](https://pre-commit.com/) to automatically run code quality checks before each commit. Pre-commit hooks help maintain consistent code style and catch common issues early.
+
+#### Installation
+
+1. Install pre-commit (if not already installed):
+
+```bash
+pip install pre-commit
+```
+
+2. Install the git hook scripts:
+
+```bash
+pre-commit install
+```
+
+#### What Gets Checked
+
+The pre-commit hooks will automatically run:
+
+- **pyink**: Google's Python code formatter (enforces 80-char line length, 2-space indentation)
+- **pylint**: Python linter following Google's style guide
+- **Standard checks**: trailing whitespace, file endings, YAML/JSON validation, merge conflicts, etc.
+
+#### Usage
+
+Once installed, the hooks run automatically on `git commit`. If any hook fails:
+
+1. Review the errors/warnings
+2. Fix the issues (some hooks auto-fix formatting)
+3. Stage the changes: `git add .`
+4. Commit again: `git commit`
+
+#### Running Manually
+
+To run pre-commit hooks on all files without committing:
+
+```bash
+pre-commit run --all-files
+```
+
+To run a specific hook:
+
+```bash
+pre-commit run pyink --all-files
+pre-commit run pylint --all-files
+```
+
+#### Skipping Hooks (Not Recommended)
+
+In rare cases where you need to skip hooks:
+
+```bash
+git commit --no-verify
+```
+
+**Note**: CI checks will still run on pull requests, so it's best to fix issues locally.
+
 ## Contribution process
 
 ### Code Reviews