format

Author: ASKabalan

tests/test_sharded_array.py

@@ -3,7 +3,6 @@
 #os.environ["JAX_PLATFORM_NAME"] = "cpu"
 #os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"
 
-
 os.environ["EQX_ON_ERROR"] = "nan"
 from functools import partial
 
@@ -25,7 +24,6 @@
 
 #assert jax.device_count() >= 8, "This notebook requires a TPU or GPU runtime with 8 devices"
 
-
 all_gather = partial(process_allgather, tiled=False)
 
 pdims = (2, 4)