Cifar dataloader does not work

The cifar dataloader no longer works properly with jax algorithms using jax.jit. I did not test to see if pytorch algorithms still work with cifar.

Description

When running jax_nadamw_full_budget.py optimizer with the cifar workload, an error is thrown which says
len(shards) = 128 but len(devices) = 8
Here is the relevant log:

I0918 18:15:28.313870 140360727609984 submission_runner.py:359] Starting training loop.
Traceback (most recent call last):
File "/algorithmic-efficiency/submission_runner.py", line 869, in
app.run(main)
File "/usr/local/lib/python3.11/site-packages/absl/app.py", line 308, in run
_run_main(main, args)
File "/usr/local/lib/python3.11/site-packages/absl/app.py", line 254, in _run_main
sys.exit(main(argv))
^^^^^^^^^^
File "/algorithmic-efficiency/submission_runner.py", line 834, in main
score = score_submission_on_workload(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/algorithmic-efficiency/submission_runner.py", line 747, in score_submission_on_workload
score, _ = train_once(
^^^^^^^^^^^
File "/algorithmic-efficiency/submission_runner.py", line 375, in train_once
batch = data_selection(
^^^^^^^^^^^^^^^
File "/algorithmic-efficiency/algorithms/baselines/self_tuning/jax_nadamw_full_budget.py", line 446, in data_selection
batch = next(input_queue)
^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/site-packages/flax/jax_utils.py", line 147, in prefetch_to_device
enqueue(size) # Fill up the buffer.
^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/site-packages/flax/jax_utils.py", line 145, in enqueue
queue.append(jax.tree_util.tree_map(_prefetch, data))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/site-packages/jax/_src/tree_util.py", line 361, in tree_map
return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/site-packages/jax/_src/tree_util.py", line 361, in
return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
^^^^^^
File "/usr/local/lib/python3.11/site-packages/flax/jax_utils.py", line 141, in _prefetch
return jax.device_put_sharded(list(xs), devices)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/site-packages/jax/_src/api.py", line 2636, in device_put_sharded
raise ValueError(f"len(shards) = {len(shards)} must equal "
ValueError: len(shards) = 128 must equal len(devices) = 8.
2025-09-18 18:15:29.477815: W tensorflow/core/kernels/data/cache_dataset_ops.cc:916] The calling iterator did not fully read the dataset being cached. In order to avoid unexpected truncation of the dataset, the partially cached contents of the dataset will be discarded. This can happen if you have an input pipeline similar to dataset.cache().take(k).repeat(). You should use dataset.take(k).cache().repeat() instead.

Steps to Reproduce

1. In algorithms/baselines/self_tuning/jax_nadamw_full_budget.py add the following two lines in get_batch_size function:

elif workload_name == 'cifar':
return 128

2. Then run the cifar workload in docker:
   python submission_runner.py
   --framework=jax
   --workload=cifar
   --experiment_dir=/experiment_runs
   --experiment_name=jax_debug_cifar
   --data_dir=/data
   --tuning_ruleset=self
   --submission_path=algorithms/baselines/self_tuning/jax_nadamw_full_budget.py

Source or Possible Fix

I think the cifar is not an officially supported workload, but it can be useful for debugging. So once it is not too much trouble we should fix this.

Start with algorithms/baselines/self_tuning/jax_nadamw_full_budget.py, especially get_batch_size and data_selection, then reproduce the failure using the submission_runner.py command and the CIFAR workload described above. Trace how the batch reaches Flax prefetch_to_device and verify the CIFAR JAX run completes without a shard/device length mismatch.