code: def test_layer_choice(
      space_format: Type[ExecutableModelSpace], nn: ModuleType
  ):
      class Net(ModelSpace, label_prefix="model"):
          def __init__(self):
              super().__init__()
              self.module = LayerChoice(
                  [
                      nn.Conv2d(3, 2, kernel_size=1),
                      nn.Conv2d(3, 5, kernel_size=1),
                  ]
              )
  
          def forward(self, x):
              return self.module(x)
  
      model = Net()
      with pytest.raises(KeyError, match="allow adding"):
          model.module[2] = nn.Conv2d(3, 7, kernel_size=1)
  
      model.module[0] = nn.Conv2d(3, 3, kernel_size=1)
  
      assert model.module[1].out_channels == 5
  
      model = space_format.from_model(model)
  
      space = model.simplify()
      assert len(space) == 1
      assert isinstance(next(iter(space.values())), Categorical)
  
      layer_label = next(iter(space.keys()))
      choices = space[layer_label].values
  
      model1 = model.freeze(
          {layer_label: choices[0]}
      ).executable_model()
      model2 = model.freeze(
          {layer_label: choices[1]}
      ).executable_model()
  
      if space_format == SimplifiedModelSpace:
          # replace doesn't work in simplified space
          assert model1(torch.randn(1, 3, 3, 3)).size() == torch.Size(
              [1, 2, 3, 3]
          )
      else:
          assert model1(torch.randn(1, 3, 3, 3)).size() == torch.Size(
              [1, 3, 3, 3]
          )
      assert model2(torch.randn(1, 3, 3, 3)).size() == torch.Size(
          [1, 5, 3, 3]
      )
  
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 4 for calling a global function with example usage: test_layer_choice(space_format, nn) and returns: self