diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index 04f035ddc..b00364908 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -204,6 +204,7 @@ jobs:
         pytest -vx tests/test_param_shapes.py
         pytest -vx tests/test_param_types.py
         pytest -vx tests/test_ssim.py
+        pytest -vx tests/submission_runner_test.py
   pytest-baselines:
     runs-on: ubuntu-latest
     steps:
diff --git a/algorithms/development_algorithms/mnist/mnist_jax/submission.py b/algorithms/development_algorithms/mnist/mnist_jax/submission.py
index afdd1bd43..8581df747 100644
--- a/algorithms/development_algorithms/mnist/mnist_jax/submission.py
+++ b/algorithms/development_algorithms/mnist/mnist_jax/submission.py
@@ -1,15 +1,12 @@
 """Training algorithm track submission functions for MNIST."""
 
-import functools
 from typing import Any, Dict, Iterator, List, Optional, Tuple
 
 import jax
 import jax.numpy as jnp
 import optax
-from flax import jax_utils
-from jax import lax
 
-from algoperf import spec
+from algoperf import jax_sharding_utils, spec
 
 
 def get_batch_size(workload_name):
@@ -39,29 +36,18 @@ def init_optimizer_state(
     ),
     optax.scale(-hyperparameters.learning_rate),
   )
-  return jax_utils.replicate(opt_init_fn(params_zeros_like)), opt_update_fn
-
-
-# We need to jax.pmap here instead of inside update_params because the latter
-# would recompile the function every step.
-@functools.partial(
-  jax.pmap,
-  axis_name='batch',
-  in_axes=(None, None, 0, 0, None, 0, 0, 0),
-  static_broadcasted_argnums=(0, 1),
-)
-def pmapped_update_params(
+  return opt_init_fn(params_zeros_like), opt_update_fn
+
+
+def train_step(
   workload: spec.Workload,
   opt_update_fn,
   current_param_container: spec.ParameterContainer,
   model_state: spec.ModelAuxiliaryState,
-  hyperparameters: spec.Hyperparameters,
   batch: Dict[str, spec.Tensor],
   optimizer_state: spec.OptimizerState,
   rng: spec.RandomState,
 ) -> spec.UpdateReturn:
-  del hyperparameters
-
   def loss_fn(params):
     logits_batch, new_model_state = workload.model_fn(
       params=params,
@@ -77,7 +63,6 @@ def loss_fn(params):
 
   grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
   (_, new_model_state), grad = grad_fn(current_param_container)
-  grad = lax.pmean(grad, axis_name='batch')
   updates, new_optimizer_state = opt_update_fn(
     grad, optimizer_state, current_param_container
   )
@@ -105,18 +90,33 @@ def update_params(
   del train_state
   del eval_results
   del global_step
+  del hyperparameters
 
-  per_device_rngs = jax.random.split(rng, jax.local_device_count())
   optimizer_state, opt_update_fn = optimizer_state
-  new_optimizer_state, updated_params, new_model_state = pmapped_update_params(
+
+  replicated = jax_sharding_utils.get_replicate_sharding()
+  sharded = jax_sharding_utils.get_batch_dim_sharding()
+  jitted_train_step = jax.jit(
+    train_step,
+    static_argnums=(0, 1),
+    donate_argnums=(2, 3, 5),
+    in_shardings=(
+      replicated,
+      replicated,
+      sharded,
+      replicated,
+      replicated,
+    ),
+    out_shardings=(replicated, replicated, replicated),
+  )
+  new_optimizer_state, updated_params, new_model_state = jitted_train_step(
     workload,
     opt_update_fn,
     current_param_container,
     model_state,
-    hyperparameters,
     batch,
     optimizer_state,
-    per_device_rngs,
+    rng,
   )
   return (new_optimizer_state, opt_update_fn), updated_params, new_model_state
 
diff --git a/pyproject.toml b/pyproject.toml
index 74e369266..36c7f8ae1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -40,6 +40,7 @@ dependencies = [
   "pandas>=2.0.1",
   "tensorflow==2.19.0",
   "tensorflow-datasets==4.9.9",
+  "tensorflow-metadata==1.16.1",
   "tensorflow-probability==0.20.0",
   "gputil==1.4.0",
   "psutil==6.1.0",
diff --git a/tests/submission_runner_test.py b/tests/submission_runner_test.py
index c6c993b7b..c68087ca5 100644
--- a/tests/submission_runner_test.py
+++ b/tests/submission_runner_test.py
@@ -14,13 +14,14 @@
 
 import submission_runner
 from algoperf.profiler import PassThroughProfiler
+from algoperf.workloads import workloads
 
 FLAGS = flags.FLAGS
 # Needed to avoid UnparsedFlagAccessError
 # (see https://github.com/google/model_search/pull/8).
-FLAGS(sys.argv)
+FLAGS([sys.argv[0]])
 
-_MNIST_DEV_ALGO_DIR = 'reference_algorithms/development_algorithms/mnist'
+_MNIST_DEV_ALGO_DIR = 'algorithms/development_algorithms/mnist'
 
 
 class SubmissionRunnerTest(parameterized.TestCase):
@@ -46,13 +47,14 @@ def test_submission(
     self, workload, framework, submission_path, tuning_search_space
   ):
     FLAGS.framework = framework
-    workload_metadata = copy.deepcopy(submission_runner.WORKLOADS[workload])
+    FLAGS.torch_compile = False
+    workload_metadata = copy.deepcopy(workloads.WORKLOADS[workload])
     workload_metadata['workload_path'] = os.path.join(
-      submission_runner.BASE_WORKLOADS_DIR,
+      workloads.BASE_WORKLOADS_DIR,
       workload_metadata['workload_path'] + '_' + framework,
       'workload.py',
     )
-    workload_obj = submission_runner.import_workload(
+    workload_obj = workloads.import_workload(
       workload_path=workload_metadata['workload_path'],
       workload_class_name=workload_metadata['workload_class_name'],
       workload_init_kwargs={},
@@ -73,7 +75,7 @@ def test_submission(
   def test_convert_filepath_to_module(self):
     """Sample test for the `convert_filepath_to_module` function."""
     test_path = os.path.abspath(__file__)
-    module_path = submission_runner.convert_filepath_to_module(test_path)
+    module_path = workloads.convert_filepath_to_module(test_path)
     self.assertNotIn('.py', module_path)
     self.assertNotIn('/', module_path)
     self.assertIsInstance(module_path, str)