# Path: rigl/experimental/jax/shuffled_mask.py
# def main(argv: List[str]):


import glob
import tempfile
from os import path
from absl.testing import absltest
from absl.testing import flagsaver
from rigl.experimental.jax import shuffled_mask

    self.assertTrue(len(files) == 1 and path.exists(files[0]))

  def test_run_random(self):
    """Test random mask driver with per-neuron sparsity."""
    experiment_dir = tempfile.mkdtemp()
    self._eval_flags = dict(
        epochs=1,
        experiment_dir=experiment_dir,
        mask_type='random',
    )

    with flagsaver.flagsaver(**self._eval_flags):
      shuffled_mask.main([])

    outfile = path.join(experiment_dir, '*', 'events.out.tfevents.*')
    files = glob.glob(outfile)

    self.assertTrue(len(files) == 1 and path.exists(files[0]))

  def test_run_per_neuron(self):
    """Test random mask driver with per-neuron sparsity."""
    experiment_dir = tempfile.mkdtemp()
    self._eval_flags = dict(
        epochs=1,
        experiment_dir=experiment_dir,
        mask_type='per_neuron',
    )

    with flagsaver.flagsaver(**self._eval_flags):

# Path: rigl/experimental/jax/shuffled_mask.py
# def main(argv: List[str]):


# Path: rigl/experimental/jax/shuffled_mask_test.py
import glob
import tempfile
from os import path
from absl.testing import absltest
from absl.testing import flagsaver
from rigl.experimental.jax import shuffled_mask


    self.assertTrue(len(files) == 1 and path.exists(files[0]))

  def test_run_random(self):
    """Test random mask driver with per-neuron sparsity."""
    experiment_dir = tempfile.mkdtemp()
    self._eval_flags = dict(
        epochs=1,
        experiment_dir=experiment_dir,
        mask_type='random',
    )

    with flagsaver.flagsaver(**self._eval_flags):
      shuffled_mask.main([])

    outfile = path.join(experiment_dir, '*', 'events.out.tfevents.*')
    files = glob.glob(outfile)

    self.assertTrue(len(files) == 1 and path.exists(files[0]))

  def test_run_per_neuron(self):
    """Test random mask driver with per-neuron sparsity."""
    experiment_dir = tempfile.mkdtemp()
    self._eval_flags = dict(
        epochs=1,
        experiment_dir=experiment_dir,
        mask_type='per_neuron',
    )

    with flagsaver.flagsaver(**self._eval_flags):

      shuffled_mask.main([])