Melissa Tan

Opening an issue for the discussion linked below. The proposal that @marcvanzee and I discussed today is: 1. Create a minimal example (e.g. MNIST) that uses gradient accumulation, likely the...

The linen_transforms_test file contains a unit test for TransformedMethodReturnValueError when a Module is returned (`test_returned_module_warning`), but doesn't yet cover the scenario where a Variable is returned.