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


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

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

  def test_prune_global_pruning_schedule(self):
    """Tests training/pruning driver with a global sparsity schedule."""
    experiment_dir = self.create_tempdir().full_path
    eval_flags = dict(
        epochs=10,
        pruning_schedule='[(5, 0.33), (7, 0.66), (9, 0.95)]',
        experiment_dir=experiment_dir,
    )

    with flagsaver.flagsaver(**eval_flags):
      prune.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_prune_local_pruning_schedule(self):
    """Tests training/pruning driver with a single layer sparsity schedule."""
    experiment_dir = self.create_tempdir().full_path
    eval_flags = dict(
        epochs=10,
        pruning_schedule='{1:[(5, 0.33), (7, 0.66), (9, 0.95)]}',
        experiment_dir=experiment_dir,
    )

    with flagsaver.flagsaver(**eval_flags):

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


# Path: rigl/experimental/jax/prune_test.py
import glob
from os import path
from absl.testing import absltest
from absl.testing import flagsaver
from rigl.experimental.jax import prune


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

  def test_prune_global_pruning_schedule(self):
    """Tests training/pruning driver with a global sparsity schedule."""
    experiment_dir = self.create_tempdir().full_path
    eval_flags = dict(
        epochs=10,
        pruning_schedule='[(5, 0.33), (7, 0.66), (9, 0.95)]',
        experiment_dir=experiment_dir,
    )

    with flagsaver.flagsaver(**eval_flags):
      prune.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_prune_local_pruning_schedule(self):
    """Tests training/pruning driver with a single layer sparsity schedule."""
    experiment_dir = self.create_tempdir().full_path
    eval_flags = dict(
        epochs=10,
        pruning_schedule='{1:[(5, 0.33), (7, 0.66), (9, 0.95)]}',
        experiment_dir=experiment_dir,
    )

    with flagsaver.flagsaver(**eval_flags):

      prune.main([])