Try pmap without a lambda to see if this allows for cache hits

Author: lukebaumann

pathwaysutils/elastic/manager.py

@@ -44,6 +44,10 @@
 _logger = logging.getLogger(__name__)
 
 
+def _plus_one(x):
+  return x + 1
+
+
 class ElasticRuntimeError(RuntimeError):
   """Error raised when elasticity cannot continue.
 
@@ -234,7 +238,7 @@ def _simple_execution(self, devices: Sequence[jax.Device]) -> jax.Array:
         self._SIMPLE_EXECUTION_TEST_VALUE - 1
     )
 
-    return jax.pmap(lambda x: x + 1, devices=devices)(test_input)
+    return jax.pmap(_plus_one, devices=devices)(test_input)
 
   @timing.timeit
   def get_slice_availability(self) -> set[int]: