Repository: google/flax

Issue:
Mutable / Unmutable state when training
Hi Flax team,

I'm working on a model with an internal state which gets updated during training. When calling the model during validation, I do not want to update these variables. I could technically introduce a training variable, but I feel it could be more elegantly (and much simpler!) by simply checking if a state is mutable:

1 ) Is there an easy way to check if a variable is mutable, so that my code only updates when it is? e.g.
`if is_initialized and var.is_mutable(): var.value = f(....)`
2 ) If I set mutable=False, I only get back the output. Is there a way to get back the state regardless? e.g.
`output, updated_state = model.apply(inputs, mutable=False, return_state=True)`
My usecase is that for my validation metrics I call my loss function with the testdata and extract the metrics, so that I can use the same code for both training and validation.

Thanks!


Discussion:
The code for this is already there it is just not exposed as an API. This will be pretty easy.

We decided not to return state that isn't updated. the reasoning is that it becomes easy to accidentally store the same variables twice or return a copy of variables from a compiled function that you don't need.

I'd be happy to give it a try and implement / write an example if you could give me some pointers; I couldn't find anything in the source code when I looked at it though...

Alright, I understand, but still a shame. Adding a keyword also goes against the design?

Resolving patch:
diff --git a/flax/core/scope.py b/flax/core/scope.py
--- a/flax/core/scope.py
+++ b/flax/core/scope.py
@@ -223,6 +223,10 @@ def value(self, value: T):
     """Updates the value of this Variable."""
     self.scope.put_variable(self.collection, self.name, value)
 
+  def is_mutable(self) -> bool:
+    """Checks if this Variable is mutable."""
+    return self.scope.is_mutable_collection(self.collection)
+
 
 class Scope:
   """A Scope allows easy access to variables and manages RNGS of a neural network layer.
diff --git a/flax/linen/module.py b/flax/linen/module.py
--- a/flax/linen/module.py
+++ b/flax/linen/module.py
@@ -654,6 +654,12 @@ def has_variable(self, col: str, name: str) -> bool:
       raise ValueError("Can't access variables on unbound modules")
     return self.scope.has_variable(col, name)
 
+  def is_mutable_collection(self, col: str) -> bool:
+    """Returns true if the collection `col` is mutable."""
+    if self.scope is None:
+      raise ValueError("Can't check mutability on unbound modules")
+    return self.scope.is_mutable_collection(col)
+
   def make_rng(self, name: str) -> PRNGKey:
     """Returns a new RNG key from a given RNG sequence for this Module.
     


Test patch:
diff --git a/tests/core/scope_test.py b/tests/core/scope_test.py
--- a/tests/core/scope_test.py
+++ b/tests/core/scope_test.py
@@ -83,6 +83,14 @@ def f(scope):
     with self.assertRaisesWithLiteralMatch(ValueError, 'No paramater named "kernel" exists in "/dense".'):
       apply(f)({})
 
+  def test_variable_is_mutable(self):
+    def f(scope, should_be_mutable):
+      test = scope.variable('state', 'test', lambda: 1)
+      self.assertEqual(test.is_mutable(), should_be_mutable)
+
+    _, variables = apply(f, mutable='state')({}, True)
+    apply(f, mutable=False)(variables, False)
+
 
 if __name__ == '__main__':
   absltest.main()
diff --git a/tests/linen/module_test.py b/tests/linen/module_test.py
--- a/tests/linen/module_test.py
+++ b/tests/linen/module_test.py
@@ -734,6 +734,15 @@ def __call__(self):
     foo = Foo()
     with self.assertRaisesWithLiteralMatch(TypeError, "Module instance is frozen outside of setup method."):
       foo.init(random.PRNGKey(0))
+  
+  def test_is_mutable_collection(self):
+    class EmptyModule(nn.Module):
+      def __call__(self):
+        return self.is_mutable_collection('test')
+
+    empty = EmptyModule()
+    self.assertTrue(empty.apply({}, mutable=['test'])[0])
+    self.assertFalse(empty.apply({}, mutable=False))
 
 
 if __name__ == '__main__':