def test_setattr(self):
mod = deferred.Deferred(ExampleModule)
mod()
new_w = jnp.ones_like(mod.w)
mod.w = new_w
self.assertIs(mod.w, new_w)
self.assertIs(mod.target.w, new_w) # pytype: disable=attribute-error
def test_only_computes_target_once(self):
target = ExampleModule()
targets = [target]
mod = deferred.Deferred(targets.pop) # pytype: disable=wrong-arg-types
for _ in range(10):
# If target was recomputed more than once pop should fail.
self.assertIs(mod.target, target)
self.assertEmpty(targets)
def test_setattr_on_target(self):
mod = deferred.Deferred(ExampleModule)
mod()
w = jnp.ones_like(mod.w)
mod.w = None
# Assigning to the target directly should reflect in the parent.
mod.target.w = w
self.assertIs(mod.w, w)
self.assertIs(mod.target.w, w)
def test_str(self):
m = ExampleModule()
d = deferred.Deferred(lambda: m)
self.assertEqual("Deferred(%s)" % m, str(d))
def test_alternative_forward_call_type_error(self):
mod = deferred.Deferred(AlternativeForwardModule,
call_methods=("forward", ))
msg = "'AlternativeForwardModule' object is not callable"
with self.assertRaisesRegex(TypeError, msg):
mod()
def test_alternative_forward(self):
mod = deferred.Deferred(AlternativeForwardModule,
call_methods=("forward", ))
self.assertEqual(mod.forward(), 42)
def test_delattr(self):
mod = deferred.Deferred(ExampleModule)
mod()
self.assertTrue(hasattr(mod.target, "w"))
del mod.w
self.assertFalse(hasattr(mod.target, "w"))
def test_getattr(self):
mod = deferred.Deferred(ExampleModule)
mod()
self.assertIs(mod.w, mod.target.w) # pytype: disable=attribute-error
def test_attr_forwarding_fails_before_construction(self):
mod = deferred.Deferred(ExampleModule)
with self.assertRaises(AttributeError):
getattr(mod, "foo")
def test_target(self):
target = ExampleModule()
mod = deferred.Deferred(lambda: target)
self.assertIs(mod.target, target)
def __init__(self, name="outer"):
super().__init__(name=name)
self.deferred = deferred.Deferred(ExampleModule)
def test_deferred_naming_name_scope(self):
with module.name_scope("foo"):
d = deferred.Deferred(ExampleModule)
mod = d.target
self.assertEqual(mod.module_name, "foo/example_module")