Update github workflow to include static analysis checks. Delete test.sh and lint.sh, and move its logic to the ci.yml file.

PiperOrigin-RevId: 859707838

Author: Ryan McKenna

.github/workflows/ci.yml

@@ -21,6 +21,37 @@ jobs:
       - uses: "actions/setup-python@v6"
         with:
             python-version: "${{ matrix.python-version }}"
-      - name: Run CI tests
-        run: bash test.sh
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install build wheel
+          pip install pylint pytype flake8 pylint-exit
+          pip install -e .[dev]
+          pip install -r docs/requirements.txt
+          pip install crc32c
+      - name: Run flake8
+        run: flake8 jax_privacy tests examples
+      - name: Run pylint
+        run: |
+          pylint --ignore=auditing.py --rcfile=.pylintrc jax_privacy || pylint-exit -efail -wfail -cfail -rfail $?
+          # pylint --rcfile=.pylintrc examples || pylint-exit -efail -wfail -cfail -rfail $?
+          # pylint --rcfile=.pylintrc tests -d W0212,C0114 || pylint-exit -efail -wfail -cfail -rfail $?
         shell: bash
+      - name: Run pytype
+        run: |
+          pytype jax_privacy -k
+          pytype tests -k
+          pytype examples -k
+      -name: Run doctests
+       run: |
+         pytest --doctest-modules jax_privacy/
+      - name: Run 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"
+          export HYPOTHESIS_PROFILE=dpftrl_default
+          pytest -n auto tests/ -k "matrix_factorization" --ignore=tests/matrix_factorization/buffered_toeplitz_test.py
+        shell: bash
+      - name: Build docs
+        run: sphinx-build -W -b html docs/ docs/_build/html

MANIFEST.in

@@ -293,7 +293,7 @@ def batch_iterator(
 
 @dataclasses.dataclass(frozen=True)
 class UserSelectionStrategy:
-  """Applies base_strategy at the user level, and selects multiple examples per user.
+  """Applies base_strategy at user level, selecting multiple examples per user.
 
   Each batch returned by the batch_iterator is a 2D array of integer indices,
   where all entries in the same row are examples owned by the same user. The

conftest.py

@@ -55,7 +55,7 @@ def __call__(self, *args, **kwargs):
 
   def sensitivity(
       self,
-      neighboring_relation: dp_accounting.NeighboringRelation = _REPLACE_SPECIAL,
+      neighboring_relation: dp_accounting.NeighboringRelation = _REPLACE_SPECIAL
   ):
     """Returns the L2 sensitivity of the Callable.
 

jax_privacy/batch_selection.py

@@ -20,19 +20,21 @@
   >>> import dp_accounting
   >>> event = dpsgd_event(noise_multiplier=3, iterations=512, sampling_prob=0.1)
   >>> accountant = dp_accounting.pld.PLDAccountant()
-  >>> accountant.compose(event).get_epsilon(target_delta=1e-6)
-  3.769686577205207
+  >>> epsilon = accountant.compose(event).get_epsilon(target_delta=1e-6)
+  >>> round(epsilon, 2)
+  3.77
 
 Example Usage (Calibrating Noise Multiplier for DP-SGD):
 
   >>> make_event = lambda sigma: dpsgd_event(sigma, 512, sampling_prob=0.1)
-  >>> dp_accounting.calibrate_dp_mechanism(
+  >>> noise_multiplier = dp_accounting.calibrate_dp_mechanism(
   ...     dp_accounting.pld.PLDAccountant,
   ...     make_event,
   ...     target_epsilon=1.0,
   ...     target_delta=1e-6
   ... )
-  9.661830157637436
+  >>> round(noise_multiplier, 2)
+  9.66
 
 Example Usage (Calibrating Number of Iterations for DP-SGD):
 
@@ -105,7 +107,7 @@ def truncated_dpsgd_event(
     num_examples: int,
     truncated_batch_size: int,
 ) -> dp_accounting.DpEvent:
-  """Returns the DpEvent for truncated DP-SGD with the given training parameters.
+  """Returns the DpEvent for truncated DP-SGD with the given training params.
 
   This mechanism is like DP-SGD, but batches larger than `truncated_batch_size`
   are truncated to size `truncated_batch_size`. See these references for more

jax_privacy/clipping.py

@@ -19,7 +19,7 @@
 
 from collections.abc import Callable
 import dataclasses
-from typing import Any, TypeAlias, TypeVar
+from typing import Any, Generic, TypeAlias, TypeVar
 
 import chex
 import jax
@@ -37,7 +37,7 @@
 
 
 @dataclasses.dataclass(frozen=True)
-class StreamingMatrix:
+class StreamingMatrix(Generic[State]):
   """A linear mapping x -> A x for a lower-triangular (streaming) A matrix.
 
   Via the attributes / member functions `init_multiply` and `multiply_next`,

jax_privacy/experimental/accounting.py

@@ -59,7 +59,7 @@
 
 
 class _NoiseStructureFn(Protocol):
-  """Protocol specifying the semantics of _IntermediateStrategy.get_noise_structure.
+  """Protocol defining semantics of _IntermediateStrategy.get_noise_structure.
 
   (Expected) Formal Guarantees of y = get_noise_structure(x):
     - x.size <= y.size (i.e., math.prod(x.shape) <= math.prod(y.shape)). See