| 493 | self.assertEqual(mod.variables[0], mod.normal_variable) |
| 494 | |
| 495 | def test_with_path(self): |
| 496 | mod = module.Module() |
| 497 | mod.w = variables.Variable(1.) |
| 498 | mod.encoder = module.Module() |
| 499 | mod.encoder.w = [({"k": mod.w}, {"k": mod.w})] |
| 500 | mod.decoder = mod.encoder |
| 501 | |
| 502 | state_dict = dict( |
| 503 | mod._flatten(with_path=True, predicate=module._is_variable)) |
| 504 | |
| 505 | self.assertEqual(state_dict, |
| 506 | {("w",): mod.w, |
| 507 | ("encoder", "w", 0, 0, "k"): mod.encoder.w[0][0]["k"], |
| 508 | ("encoder", "w", 0, 1, "k"): mod.encoder.w[0][1]["k"], |
| 509 | ("decoder", "w", 0, 0, "k"): mod.decoder.w[0][0]["k"], |
| 510 | ("decoder", "w", 0, 1, "k"): mod.decoder.w[0][1]["k"]},) |
| 511 | |
| 512 | def test_module_discover_layer_variable(self): |
| 513 | m = module.Module() |