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()
