fix: Optimize Batch Selection and Update JAX Compatibility

Author: amyssnippet

examples/distributed_noise_generation.py

@@ -36,6 +36,8 @@
 from absl import flags
 import jax
 import jax.numpy as jnp
+# pylint: disable=g-importing-member
+from jax.sharding import reshard
 from jax_privacy import noise_addition
 from jax_privacy.matrix_factorization import toeplitz
 
@@ -162,10 +164,10 @@ def run(pytree_like_model_params):
   t0 = time.time()
   compiled_run = run.lower(model_params).compile()
   t1 = time.time()
-  print('[BandMF] Compilation time: %.3f seconds' % (t1 - t0))
+  print(f'[BandMF] Compilation time: {t1 - t0:.3f} seconds')
   state, noisy_grad = jax.block_until_ready(compiled_run(model_params))
   t2 = time.time()
-  print('[BandMF] Per-step run time: %.3f seconds' % ((t2 - t1) / steps))
+  print(f'[BandMF] Per-step run time: {(t2 - t1) / steps:.3f} seconds')
 
   return state, noisy_grad
 

jax_privacy/batch_selection.py

@@ -293,7 +293,9 @@ def batch_iterator(
 
 @dataclasses.dataclass(frozen=True)
 class UserSelectionStrategy:
-  """A strategy that applies a base_strategy at the user level.
+  """Applies base_strategy at the user level, and selects 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
@@ -354,12 +356,19 @@ def batch_iterator(
     num_examples = user_ids.size
     dtype = np.min_scalar_type(-num_examples)
 
+    # Precompute sorted indices and starts to avoid O(n) per user_id
+    # in np.where.
+    sorted_indices = np.argsort(inverse)
+    counts = np.bincount(inverse, minlength=num_users)
+    starts = np.r_[0, np.cumsum(counts)]
+
     def create_user_generator(user_id):
-      # TODO: b/415360727 - this where is suboptimal, as it is O(n) per user_id.
-      owned_examples = np.where(inverse == user_id)[0].astype(dtype)
+      start = starts[user_id]
+      end = starts[user_id + 1]
+      owned_examples = sorted_indices[start:end].astype(dtype)
       if self.shuffle_per_user:
         rng.shuffle(owned_examples)
-      return itertools.cycle(list(owned_examples))
+      return itertools.cycle(owned_examples)
 
     user_generators = [create_user_generator(i) for i in range(num_users)]
 

jax_privacy/sharding_utils.py

@@ -43,7 +43,9 @@ def _ceiling_to_multiple(size: int, multiple: int) -> int:
   return size + multiple - remainder if remainder != 0 else size
 
 
-def flatten_with_zero_redundancy(abstract_array: jax.Array) -> jax.Array:
+def flatten_with_zero_redundancy(
+    abstract_array: jax.ShapeDtypeStruct | jax.Array
+) -> jax.ShapeDtypeStruct:
   """Return a flattened, padded, and ZeRo-sharded abstract version of x.
 
   Specifically, the returned object will describe a 1D array that is
@@ -59,15 +61,11 @@ def flatten_with_zero_redundancy(abstract_array: jax.Array) -> jax.Array:
     A zero-redundancy abstract flattened+padded version of the input value.
   """
   mesh = jax.typeof(abstract_array).sharding.mesh
-  # As of JAX 0.7.0, jnp.*_like will not preserve sharding of ShapeDtypeStruct
-  # defined w.r.t. AbstractMeshes, so we return a concrete array here.
-  # Under JIT, this should get optimized away.
-  # TODO: b/415360727 - Version bump to 0.7.1, swap in jax.ShapeDtypeStruct,
-  # and add type annotations to this function.
-  return jax.numpy.empty(
-      _ceiling_to_multiple(abstract_array.size, mesh.size),
+  # As of JAX 0.7.1, we can use ShapeDtypeStruct with sharding preserved.
+  return jax.ShapeDtypeStruct(
+      shape=(_ceiling_to_multiple(abstract_array.size, mesh.size),),
       dtype=abstract_array.dtype,
-      out_sharding=jax.sharding.NamedSharding(mesh, jax.P(mesh.axis_names)),
+      sharding=jax.sharding.NamedSharding(mesh, jax.P(mesh.axis_names)),
   )
 
 

pyproject.toml

@@ -14,8 +14,8 @@ requires-python = ">=3.11"
 dependencies = [
     "absl-py",
     "dp_accounting @ git+https://github.com/google/differential-privacy.git#subdirectory=python/dp_accounting",
-    "jax>=0.7.0",
-    "jaxlib>=0.7.0",
+    "jax>=0.7.1",
+    "jaxlib>=0.7.1",
     "pydantic",
     "numpy",
     "optax",

tests/sharding_utils_test.py

@@ -80,6 +80,20 @@ def test_ceiling_to_multiple(self):
     self.assertEqual(sharding_utils._ceiling_to_multiple(4, 4), 4)
     self.assertEqual(sharding_utils._ceiling_to_multiple(5, 4), 8)
 
+  def test_flatten_zeros_like_preserves_metadata(self):
+    sharding = jax.sharding.NamedSharding(
+        self.mesh, jax.sharding.PartitionSpec(None, 'y')
+    )
+    x = jax.device_put(jnp.ones((2, 6), dtype=jnp.float32), sharding)
+    flattened = sharding_utils.flatten_with_zero_redundancy(x)
+    zeros = jnp.zeros_like(flattened)
+    self.assertEqual(zeros.shape, flattened.shape)
+    self.assertEqual(zeros.dtype, flattened.dtype)
+    self.assertEqual(zeros.sharding.spec, flattened.sharding.spec)
+    self.assertEqual(
+        zeros.sharding.mesh.axis_names, flattened.sharding.mesh.axis_names
+    )
+
 
 if __name__ == '__main__':
   absltest.main()