def setUp(self): super().setUp() self._original_config = {} for key, value in self._default_config.items(): self._original_config[key] = config._read(key) config.update(key, value) # We use the adler32 hash for two reasons. # a) it is deterministic run to run, unlike hash() which is randomized. # b) it returns values in int32 range, which RandomState requires. self._rng = npr.RandomState(zlib.adler32(self._testMethodName.encode()))
def restore_spmd_manual_lowering_flag(): if old_spmd_manual_lowering_flag is None: return config.update('experimental_xmap_spmd_lowering_manual', old_spmd_manual_lowering_flag)
def set_spmd_manual_lowering_flag(val: bool): global old_spmd_manual_lowering_flag old_spmd_manual_lowering_flag = config.experimental_xmap_spmd_lowering_manual config.update('experimental_xmap_spmd_lowering_manual', val)
def tearDown(self): for key, value in self._original_config.items(): config.update(key, value) super().tearDown()