Merge pull request #2729 from levskaya:queue_runner_fix

PiperOrigin-RevId: 495767913

Author: Flax Authors

.github/workflows/build.yml

@@ -79,7 +79,7 @@ jobs:
   tests:
     name: Run Tests
     needs: [cancel-previous, pre-commit, commit-count, test-import]
-    runs-on: ubuntu-20.04-16core
+    runs-on: ubuntu-latest
     strategy:
       matrix:
         python-version: ['3.8', '3.9', '3.10']

tests/checkpoints_test.py

@@ -374,30 +374,6 @@ def test_auto_restore(self):
         os.path.join(tmp_dir, 'test_1'), target=target)
     check_eq(restored, to_save)
 
-
-  # This is for fully addressable JAX arrays. For multiprocess JAX arrays like
-  # GDA, see multihost_test.py (internal only)
-  @parameterized.parameters({'use_orbax': True, 'jax_array_config': True},
-                            {'use_orbax': False, 'jax_array_config': False})
-  def test_jax_array(self, use_orbax, jax_array_config):
-    config.flax_use_orbax_checkpointing = use_orbax
-    jax.config.update('jax_array', jax_array_config)
-    tmp_dir = pathlib.Path(self.create_tempdir().full_path)
-    test_object0 = {'a': jnp.zeros(3), 'b': jnp.arange(3)}
-    test_object1 = {'a': jnp.ones(3), 'b': jnp.arange(3, 6)}
-    new_object = checkpoints.restore_checkpoint(
-        tmp_dir, test_object0, prefix='test_')
-    check_eq(new_object, test_object0)
-    # Create leftover temporary checkpoint, which should be ignored.
-    io.GFile(os.path.join(tmp_dir, 'test_tmp'), 'w')
-    checkpoints.save_checkpoint(
-        tmp_dir, test_object1, 0, prefix='test_', keep=1)
-    self.assertIn('test_0', os.listdir(tmp_dir))
-    new_object = checkpoints.restore_checkpoint(
-        tmp_dir, test_object0, prefix='test_')
-    check_eq(new_object, {'a': np.ones(3), 'b': np.arange(3, 6)})
-
-
   def test_convert_pre_linen(self):
     params = checkpoints.convert_pre_linen({
         'mod_0': {