Make mypy happy

skye · May 19, 2020 · e87d0e7 · e87d0e7
1 parent 82f283a
commit e87d0e7
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/jax/interpreters/sharded_jit.py b/jax/interpreters/sharded_jit.py
@@ -93,6 +93,7 @@ def _sharded_callable(
     raise ValueError("sharded_jit only works on TPU!")
 
   num_partitions = pxla.reconcile_num_partitions(jaxpr, num_partitions)
+  assert num_partitions is not None
   if num_partitions > xb.local_device_count():
     raise ValueError(
         f"sharded_jit computation requires {num_partitions} devices, "