code: def test_ensure_frozen():
      discrete = Categorical([1, 2, 3])
      with pytest.raises(RuntimeError, match="No frozen context"):
          ensure_frozen(discrete)
  
      class Net(ModelSpace):
          def __init__(self):
              super().__init__()
              self.module = Categorical([1, 2, 3], label="x")
              self.add_mutable(self.module)
              self.dry_run_value = ensure_frozen(self.module)
  
      model = Net()
      assert model.dry_run_value == 1
  
      with pytest.raises(
          NotImplementedError, match="nni.nas.space.model_context"
      ):
          model.freeze({"x": 2})
  
imports: import copy
  import pytest
  import torch
  import nni
  from torch import nn
  from nni.common.serializer import is_traceable
  from nni.mutable import Categorical, ExpressionConstraint, ensure_frozen, frozen_context, label_scope, auto_label, SampleMissingError
  from nni.mutable.mutable import _mutable_equal
  from nni.nas.nn.pytorch import ModelSpace, MutableModule, ParametrizedModule
  from nni.nas.space import model_context
label_desc: Use image node_id 2 for calling a global function with example usage: test_ensure_frozen() without return types