code: def test_mutable_in_nn_parameter(
      space_format: Type[ExecutableModelSpace], nn: ModuleType
  ):
      class Inner(nn.Module):
          def __init__(self):
              super().__init__()
              self.aux = torch.nn.Parameter(
                  torch.zeros(
                      1,
                      ensure_frozen(
                          Categorical([64, 128, 256], label="a")
                      ),
                      3,
                      3,
                  )
              )
  
          def forward(self):
              return self.aux
  
          def freeze(self, sample):
              with model_context(sample):
                  return self.__class__()
  
      class Net(ModelSpace, label_prefix="model"):
          def __init__(self):
              super().__init__()
              self.choice = self.add_mutable(
                  Categorical([64, 128, 256], label="a")
              )
              self.inner = Inner()
              self.dim = ensure_frozen(self.choice)
  
          def forward(self):
              param = self.inner()
              assert param.size(1) == self.dim
              return param
  
          def freeze(self, sample):
              with model_context(sample):
                  return self.__class__()
  
      if issubclass(space_format, GraphModelSpace):
          with pytest.raises(RuntimeError, match="Arbitrary"):
              space_format.from_model(Net())
      else:
          model = space_format.from_model(Net())
          assert len(model.simplify()) == 1
          label = list(model.simplify())[0]
          model1 = model.freeze({label: 64})
          model2 = model.freeze({label: 256})
          assert model1.executable_model()().size(1) == 64
          assert model2.executable_model()().size(1) == 256
  
imports: import copy
  import random
  from types import ModuleType
  from typing import Type
  import pytest
  import nni
  import nni.nas.nn.pytorch.layers
  import nni.nas.evaluator.pytorch.lightning
  import torch
  import torch.nn.functional
  from nni.mutable import Categorical, ensure_frozen, label_scope, SampleValidationError
  from nni.nas.evaluator import FunctionalEvaluator
  from nni.nas.space import RawFormatModelSpace, SimplifiedModelSpace, ExecutableModelSpace, GraphModelSpace, model_context
  from nni.nas.space.pytorch import PytorchGraphModelSpace
  from nni.nas.nn.pytorch import ParametrizedModule, MutableModule, ModelSpace, LayerChoice, InputChoice, Cell, ChosenInputs, MutableConv2d, MutableLinear
label_desc: Use image node_id 25 for calling a global function with example usage: test_mutable_in_nn_parameter(space_format, nn) and returns: self, param, self, self