Merge branch 'main' into linear-gramian-computer

Author: ValerianRey

tests/unit/autogram/test_engine.py

@@ -17,6 +17,7 @@
     InstanceNormMobileNetV2,
     InstanceNormResNet18,
     InterModuleParamReuse,
+    IntraModuleParamReuse,
     MIMOBranched,
     MISOBranched,
     ModelAlsoUsingSubmoduleParamsDirectly,
@@ -40,7 +41,6 @@
     Randomness,
     RequiresGradOfSchrodinger,
     SimpleBranched,
-    SimpleParamReuse,
     SingleInputPyTreeOutput,
     SIPOBranched,
     SomeFrozenParam,
@@ -99,7 +99,7 @@
     (ModuleFactory(PIPOBranched), 32),
     (ModuleFactory(WithNoTensorOutput), 32),
     (ModuleFactory(WithBuffered), 32),
-    (ModuleFactory(SimpleParamReuse), 32),
+    (ModuleFactory(IntraModuleParamReuse), 32),
     (ModuleFactory(ModuleReuse), 32),
     (ModuleFactory(SomeUnusedParam), 32),
     (ModuleFactory(SomeFrozenParam), 32),

tests/utils/architectures.py

@@ -423,7 +423,7 @@ def forward(self, input: Tensor) -> Tensor:
         return self.linear(input)
 
 
-class SimpleParamReuse(ShapedModule):
+class IntraModuleParamReuse(ShapedModule):
     """Module that reuses the same nn.Parameter for two computations directly inside of it."""
 
     INPUT_SHAPES = (50,)