factor out simulate function

nets/experiments/in_context_learning/launcher_local.py

@@ -9,15 +9,15 @@
 import optax
 
 import nets
+from nets import datasets
+from nets import samplers
 from nets.launch import configs
 from nets.launch import submit
 from nets.launch.hparams import Param
 from nets.launch.hparams import EnumParam
 from nets.launch.hparams import FixedParam
 
-
-from nets import datasets
-from nets import samplers
+from nets.simulators.in_context_learning import simulate
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -132,6 +132,7 @@ class DebugSearchConfig(SearchConfig):
 
   jobs = submit.submit_jobs(
     executor=executor,
+    fn=simulate,
     cfg=DebugSearchConfig(
       key=jax.random.PRNGKey(0),
       num_configs=1,

nets/experiments/in_context_learning/launcher_slurm.py

@@ -9,6 +9,8 @@
 import optax
 
 import nets
+from nets import datasets
+from nets import samplers
 from nets.launch import configs
 from nets.launch import submit
 from nets.launch.hparams import Param
@@ -17,8 +19,7 @@
 from nets.launch.hparams import LogUniformParam
 from nets.launch.hparams import UniformParam
 
-from nets import datasets
-from nets import samplers
+from nets.simulators.in_context_learning import simulate
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -150,4 +151,4 @@ class SymbolicSearchConfig(SearchConfig):
     exemplar_noise_scale=LogUniformParam(1e-1, 1e3, base=10),
   )
 
-  jobs = submit.submit_jobs(executor, cfg)
+  jobs = submit.submit_jobs(executor, simulate, cfg)

nets/launch/submit.py

@@ -16,11 +16,9 @@
 
 import submitit
 
-from nets.simulators.in_context_learning import simulate
 from nets.launch import analyze
 from nets.launch import configs
 
-
 # Ignore warnings about invalid column names for PyTables.
 import warnings
 from tables import NaturalNameWarning
@@ -29,7 +27,7 @@
 
 
 def augment_df_with_kwargs(func):
-  """Return a function augments a `pd.DataFrame` with keyword arguments."""
+  """Return a function that augments a `pd.DataFrame` with keyword arguments."""
 
   def wrapped(**kwargs):
     results_df = func(**kwargs)
@@ -42,9 +40,7 @@ def wrapped(**kwargs):
 class Executor(submitit.AutoExecutor):
   """A `submitit.AutoExecutor` with a custom `starmap_array` method."""
 
-  def starmap_array(
-    self, fn: Callable, iterable: Iterable | configs.Config
-  ) -> list[Any]:
+  def starmap_array(self, fn: Callable, iterable: Iterable) -> list[Any]:
     """A distributed equivalent of the `itertools.starmap` function."""
     submissions = [
       submitit.core.utils.DelayedSubmission(fn, **kwargs) for kwargs in iterable
@@ -117,14 +113,14 @@ def get_submitit_executor(
   return executor
 
 
-def submit_jobs(executor: Executor, cfg: configs.Config):
+def submit_jobs(executor: Executor, fn: Callable, cfg: configs.Config):
   """Submit jobs to the cluster."""
   logging.info(f"Using config {pprint.pformat(cfg)}.")
 
   # Launch jobs.
   logging.info("Launching jobs...")
   jobs = executor.starmap_array(
-    augment_df_with_kwargs(simulate),
+    augment_df_with_kwargs(fn),
     cfg,
   )
   logging.info(f"Waiting for {len(jobs)} jobs to terminate...")