(self, device, dtype, module_info, training)
| 46 | |
| 47 | @modules(module_db) |
| 48 | def test_forward(self, device, dtype, module_info, training): |
| 49 | module_cls = module_info.module_cls |
| 50 | module_inputs = module_info.module_inputs_func(module_info, device=device, dtype=dtype, |
| 51 | requires_grad=False, training=training) |
| 52 | dtype_to_method_caller = { |
| 53 | torch.float32: methodcaller("float"), |
| 54 | torch.float64: methodcaller("double"), |
| 55 | } |
| 56 | for module_input in module_inputs: |
| 57 | if module_input.forward_input is None: |
| 58 | continue |
| 59 | |
| 60 | with freeze_rng_state(): |
| 61 | # === Instantiate the module. === |
| 62 | args, kwargs = module_input.constructor_input.args, module_input.constructor_input.kwargs |
| 63 | m = module_cls(*args, **kwargs) |
| 64 | m.to(device).to(dtype) |
| 65 | m.train(training) |
| 66 | |
| 67 | # === Do forward pass. === |
| 68 | args, kwargs = module_input.forward_input.args, module_input.forward_input.kwargs |
| 69 | outputs = m(*args, **kwargs) |
| 70 | |
| 71 | # === Compare outputs to a reference if one is specified. === |
| 72 | # TODO: Handle precision |
| 73 | reference_fn = module_input.reference_fn |
| 74 | if reference_fn is not None: |
| 75 | ref_outputs = reference_fn(m, *args, **kwargs) |
| 76 | self.assertEqual(outputs, ref_outputs) |
| 77 | |
| 78 | # === Use the method call and verify the parameters and buffers === |
| 79 | if dtype in dtype_to_method_caller: |
| 80 | dtype_to_method_caller[dtype](m) |
| 81 | m(*args, **kwargs) |
| 82 | self._assert_module_parameters_and_buffer_are(m, device, dtype) |
| 83 | |
| 84 | # Tests passing factory kwargs (e.g. device / dtype) during module instantiation. |
| 85 | # They should be applied to any created parameters and buffers. |
nothing calls this directly
no test coverage detected