Merge branch 'main' into feat/nnx-transformer-dp-sgd

Author: debanganghosh08

.github/workflows/ci.yml

@@ -7,22 +7,110 @@ on:
     branches: ["main"]
 
 jobs:
-  build-and-test:
-    name: "Python ${{ matrix.python-version }} on ${{ matrix.os }}"
-    runs-on: "${{ matrix.os }}"
+  lint:
+    name: "Lint: Python ${{ matrix.python-version }} on ${{ matrix.os }}"
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        python-version: ["3.11"]
+        os: [ubuntu-latest]
+    steps:
+      - uses: "actions/checkout@v6"
+      - uses: "actions/setup-python@v6"
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: "pip"
+      - name: Install Linting Tools
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e ".[lint,examples,test]"
+      - name: Run flake
+        run: |
+          flake8 jax_privacy tests examples
+      - name: Run pydocstyle
+        run: |
+          pydocstyle --convention=google --add-ignore=D101,D102,D103,D105,D202,D402 jax_privacy/
+      - name: Run pylint
+        shell: bash
+        run: |
+          pylint jax_privacy || pylint-exit -efail -wfail -cfail -rfail $?
+          pylint examples || pylint-exit -efail -wfail -cfail -rfail $?
+          pylint tests -d W0101,W0212,C0114 || pylint-exit -efail -wfail -cfail -rfail $?
+      - name: Run pytype
+        run: |
+          pytype jax_privacy -k
+          pytype tests -k
+          pytype examples -k
 
+  test:
+    name: "Unit Test: Python ${{ matrix.python-version }} on ${{ matrix.os }}"
+    runs-on: ${{ matrix.os }}
     strategy:
       matrix:
         python-version: ["3.11"]
         os: [ubuntu-latest]
+    steps:
+      - uses: "actions/checkout@v6"
+      - uses: "actions/setup-python@v6"
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: "pip"
+      - name: Install Test Dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e ".[test,examples]"
+      - name: Run Doctests
+        run: |
+          pytest --doctest-modules jax_privacy
+      - name: Run Standard Tests
+        run: |
+          pytest -n auto tests/ -k "not matrix_factorization and not distributed_noise_generation_test and not sharding_utils_test"
+          pytest -n auto tests/ -k "distributed_noise_generation_test"
+          pytest -n auto tests/ -k "sharding_utils_test"
 
+  matrix-tests:
+    name: "Matrix Tests: Python ${{ matrix.python-version }} on ${{ matrix.os }}"
+    needs: lint
+    runs-on: "${{ matrix.os }}"
+    strategy:
+      matrix:
+        python-version: ["3.11"]
+        os: [ubuntu-latest]
     steps:
       - uses: "actions/checkout@v6"
       - uses: "actions/setup-python@v6"
         with:
-            python-version: "${{ matrix.python-version }}"
-      - name: Install example requirements
-        run: pip install -r examples/requirements.txt
-      - name: Run CI tests
-        run: bash test.sh
+          python-version: "${{ matrix.python-version }}"
+          cache: "pip"
+      - name: Install Dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e ".[test]"
+      - name: Run Heavy Tests
+        run: |
+          export HYPOTHESIS_PROFILE=dpftrl_default
+          pytest -n auto tests/ -k "matrix_factorization" --ignore=tests/matrix_factorization/buffered_toeplitz_test.py
         shell: bash
+
+  docs:
+    name: "Docs: Python ${{ matrix.python-version }} on ${{ matrix.os }}"
+    needs: lint
+    runs-on: "${{ matrix.os }}"
+    strategy:
+      matrix:
+        python-version: ["3.11"]
+        os: [ubuntu-latest]
+    steps:
+      - uses: "actions/checkout@v6"
+      - uses: "actions/setup-python@v6"
+        with:
+          python-version: "${{ matrix.python-version }}"
+          cache: "pip"
+      - name: Install Docs Dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e ".[docs]"
+      - name: Build Sphinx
+        run: |
+          cd docs
+          sphinx-build -W -b html . _build/html

.pylintrc

@@ -35,6 +35,12 @@ unsafe-load-any-extension=no
 extension-pkg-whitelist=
 
 
+[STRING]
+
+# This flag controls whether inconsistent-quotes generates a warning when the
+# character used as a quote delimiter is used inconsistently within a module.
+check-quote-consistency=yes
+
 [MESSAGES CONTROL]
 
 # Only show warnings with the listed confidence levels. Leave empty to show
@@ -128,6 +134,7 @@ disable=apply-builtin,
         too-many-branches,
         too-many-instance-attributes,
         too-many-locals,
+        too-many-positional-arguments,
         too-many-public-methods,
         too-many-return-statements,
         too-many-statements,
@@ -142,6 +149,7 @@ disable=apply-builtin,
         xrange-builtin,
         wrong-import-order,
         zip-builtin-not-iterating,
+        g-bad-todo,
 
 
 [REPORTS]

CONTRIBUTING.md

@@ -3,6 +3,27 @@
 We welcome small patches related to bug fixes, but we do not plan to accept
 major changes to this repository at this time.
 
+## Coordination & Claiming Issues
+
+To avoid duplicate effort and ensure your work can be merged, please follow
+these steps:
+
+*   **Check for existing PRs and Issues:** Before starting work, search the
+    issue tracker and active PRs to see if the feature/bug is already being
+    addressed.
+
+*   **Claim the issue:** If an issue exists, comment on it expressing your
+    intent to work on it (e.g., "I'd like to work on this"). A maintainer will
+    then assign it to you.
+
+*   **Wait for Assignment:** Do not start a large-scale implementation until a
+    maintainer has acknowledged your comment. This prevents multiple people from
+    working on the same fix simultaneously.
+
+*   **Stale Assignments:** If an assigned issue hasn't seen progress or
+    communication for 14 days, the assignment may be cleared to allow others to
+    contribute.
+
 ## Contributor License Agreement
 
 Contributions to this project must be accompanied by a Contributor License
@@ -25,45 +46,47 @@ information on using pull requests.
 ## Style guide
 
 Code in this library generally follows the
-[Google Style Guide](https://google.github.io/styleguide/pyguide.html).
-We aim to keep APIs, names, and design patterns in line with the broader JAX
-ecosystem as much as possible, with immutability and functional purity being a
-key guiding principles we adhere to across our library. Below are some more
-detailed conventions depending on what is being contributed.
-
-1. **Public Facing Functions**: Public facing functions are those that are
-exposed to the **users** of JAX Privacy (usually surfaced via \_\_init\_\_.py).
-Public facing functions and classes should **always** have full docstrings, type
-annotations, and example usages in the form of
-[doctests](https://docs.python.org/3/library/doctest.html). Doctests
-provides useful documentation that stays up-to-date with code changes, and
-is a useful litmus test on the simplicity and usability of the API surface.
-
-1. **Internal functions used across files**: For maintainability of the core
-library, it is sometimes beneficial to define a function in one file and have it
-be used by another file within the jax_privacy package. These functions are not
-intended to be consumed by JAX Privacy users (although may be encountered by
-developers / contributors). These functions should generally have
-descriptive names, type annotations. Internal functions should have a one-line
-docstring explaining what they do. A full docstring is encouraged if the
-function has non-obvious side effects, complex arguments,
-or implements a multi-step algorithm that isn't clear from the code alone.
-
-1. **File-local private functions**: These function should always have a
-leading "_". This signals to developers that the function is not part of the
-public API and is subject to change without notice. These functions should have
-1-line docstrings; type annotations are optional and context-dependent. Example
-usages are not needed as they can be found in the corresponding _test.py file.
-
-1. **Nested functions**: Functions defined within other functions should
-generally be as simple as possible; we prefer to keep the boilerplate minimal
-on these (no docstrings + type annotations), inline comments can be used, but
-should be used sparingly.
+[Google Style Guide](https://google.github.io/styleguide/pyguide.html). We aim
+to keep APIs, names, and design patterns in line with the broader JAX ecosystem
+as much as possible, with immutability and functional purity being a key guiding
+principles we adhere to across our library. Below are some more detailed
+conventions depending on what is being contributed.
+
+1.  **Public Facing Functions**: Public facing functions are those that are
+    exposed to the **users** of JAX Privacy (usually surfaced via
+    \_\_init\_\_.py). Public facing functions and classes should **always** have
+    full docstrings, type annotations, and example usages in the form of
+    [doctests](https://docs.python.org/3/library/doctest.html). Doctests
+    provides useful documentation that stays up-to-date with code changes, and
+    is a useful litmus test on the simplicity and usability of the API surface.
+
+1.  **Internal functions used across files**: For maintainability of the core
+    library, it is sometimes beneficial to define a function in one file and
+    have it be used by another file within the jax_privacy package. These
+    functions are not intended to be consumed by JAX Privacy users (although may
+    be encountered by developers / contributors). These functions should
+    generally have descriptive names, type annotations. Internal functions
+    should have a one-line docstring explaining what they do. A full docstring
+    is encouraged if the function has non-obvious side effects, complex
+    arguments, or implements a multi-step algorithm that isn't clear from the
+    code alone.
+
+1.  **File-local private functions**: These function should always have a
+    leading "_". This signals to developers that the function is not part of the
+    public API and is subject to change without notice. These functions should
+    have 1-line docstrings; type annotations are optional and context-dependent.
+    Example usages are not needed as they can be found in the corresponding
+    _test.py file.
+
+1.  **Nested functions**: Functions defined within other functions should
+    generally be as simple as possible; we prefer to keep the boilerplate
+    minimal on these (no docstrings + type annotations), inline comments can be
+    used, but should be used sparingly.
 
 ## Linting and testing
 
-We use `flake8`, `pylint` and `pytype` for linting and type checking. Please
-run the following commands locally before submitting a pull request:
+We use `flake8`, `pylint` and `pytype` for linting and type checking. Please run
+the following commands locally before submitting a pull request:
 
 ```bash
 $ flake8 jax_privacy/**.py

MANIFEST.in

@@ -25,12 +25,3 @@ Experimental Modules
 
   experimental.execution_plan
   experimental.compilation_utils
-
-
-Other References
-----------------
-.. autosummary::
-  :toctree: _autosummary_output
-  :nosignatures:
-
-  experimental.microbatching

conftest.py

@@ -36,8 +36,6 @@
 from absl import flags
 import jax
 import jax.numpy as jnp
-# pylint: disable=g-importing-member,no-name-in-module
-from jax.experimental.shard import reshard
 from jax_privacy import noise_addition
 from jax_privacy.matrix_factorization import toeplitz
 
@@ -112,15 +110,17 @@ def bert_model_params(hidden_size: int) -> Any:
       is_leaf=lambda x: isinstance(x, tuple),
   )
 
-  return reshard(model_params, jax.sharding.PartitionSpec())
+  return jax.sharding.reshard(model_params, jax.sharding.PartitionSpec())
 
 
 def toy_model_params(hidden_size: int) -> jax.Array:
   """Returns model parameters for a toy model."""
   # This is a toy example where the model is just a 2D array of size (H, H^2).
   leaf_shape = (hidden_size, hidden_size**2)
 
-  return reshard(jnp.zeros(leaf_shape), jax.sharding.PartitionSpec('x', 'y'))
+  return jax.sharding.reshard(
+      jnp.zeros(leaf_shape), jax.sharding.PartitionSpec('x', 'y')
+  )
 
 
 def generate_noise(
@@ -162,10 +162,10 @@ def run(pytree_like_model_params):
   t0 = time.time()
   compiled_run = run.lower(model_params).compile()
   t1 = time.time()
-  print(f'[BandMF] Compilation time: {t1 - t0:.3f} seconds')
+  print(f'[BandMF] Compilation time: {t1-t0:.3f} seconds')
   state, noisy_grad = jax.block_until_ready(compiled_run(model_params))
   t2 = time.time()
-  print(f'[BandMF] Per-step run time: {(t2 - t1) / steps:.3f} seconds')
+  print(f'[BandMF] Per-step run time: {(t2-t1)/steps:.3f} seconds')
 
   return state, noisy_grad
 

docs/core_library.rst

@@ -22,8 +22,9 @@
 import os
 from absl import app
 
-os.environ["KERAS_BACKEND"] = "jax"  # pylint: disable=wrong-import-position
-from jax_privacy.keras import keras_api  # pylint: disable=g-import-not-at-top
+os.environ["KERAS_BACKEND"] = "jax"
+# pylint: disable=g-import-not-at-top,wrong-import-position
+from jax_privacy import keras_api
 import keras
 from keras import layers
 import numpy as np