def test_default_config_round_trip(self):
tempdir = tempfile.mkdtemp()
filepath = os.path.join(tempdir, 'config.json')
config = config_lib.get_config()
config_utils.save_config(config, filepath)
restored_config = config_utils.load_config(filepath)
self.assertTrue(config_utils.equals(config, restored_config))
def test_dataset_standard_batching(self):
dataset_name = 'control_flow_programs/decimal-L10'
data_dir = tempfile.mkdtemp()
config = config_lib.get_config()
with config.unlocked():
config.dataset.name = dataset_name
config.dataset.in_memory = True
config.dataset.batch_size = 5
config.dataset.representation = 'trace'
config = ml_collections.FrozenConfigDict(config)
dataset_info = dataset_utils.get_dataset(data_dir, config)
item = next(iter(dataset_info.dataset))
self.assertEqual(item['cfg']['data'].shape[0], 5)
def test_valid_overrides(self):
config = config_lib.get_config()
all_overrides = overrides_lib.get_all_overrides()
all_override_names = all_overrides.keys()
overrides_lib.apply_overrides(config,
override_names=all_override_names)