code: def test_contiguify_roundtrip():
      """
      Validate that contiguify round trips and reconstructions are correct.
      """
      model = DummyInferenceModel()
  
      n_layers = 2
      transformer_params = []
      transformer_containers = []
  
      # Create parameters and populate them into the containers
      for i in range(n_layers):
          transformer_containers.append(
              TransformerLayerContainer(model)
          )
          layer_params = []
          for j in range(2):
              layer_params.append(torch.rand(16, 16))
              transformer_containers[i].set_dependency(
                  f"param_{j+1}", layer_params[j]
              )
  
          layer_params = [
              p.to(get_accelerator().current_device())
              for p in layer_params
          ]
  
          transformer_params.append(layer_params)
          assert transformer_containers[i].is_populated == True
  
      non_transformer_params = []
      non_transformer_container = NonTransformerContainer(model)
  
      for i in range(3):
          non_transformer_params.append(
              torch.rand(16, 16).permute(1, 0)
          )
          non_transformer_container.set_dependency(
              f"param_{i+1}", non_transformer_params[i]
          )
  
      non_transformer_params = [
          p.to(get_accelerator().current_device())
          for p in non_transformer_params
      ]
  
      def validate_containers(
          t_containers: List[LayerContainer],
          n_t_containers: LayerContainer,
          t_params: List[List[torch.Tensor]],
          n_t_params: List[torch.Tensor],
      ):
          """
          Validate params match what is on the containers.
          """
          for i in range(n_layers):
              l_c = t_containers[i]
  
              assert l_c.is_initialized == True
  
              assert torch.equal(l_c.param_1, t_params[i][0])
              assert torch.equal(l_c.param_2, t_params[i][1])
  
          assert n_t_containers.is_initialized == True
          assert torch.equal(n_t_containers.param_1, n_t_params[0])
          assert torch.equal(n_t_containers.param_2, n_t_params[1])
          assert torch.equal(n_t_containers.param_3, n_t_params[2])
          assert not n_t_containers.param_1.is_contiguous()
          assert not n_t_containers.param_2.is_contiguous()
          assert not n_t_containers.param_3.is_contiguous()
  
      buffer, metadata = flatten_inference_model(
          transformer_containers,
          non_transformer_container,
          "NoOpPolicy",
      )
  
      # Validate containers before contiguify
      validate_containers(
          transformer_containers,
          non_transformer_container,
          transformer_params,
          non_transformer_params,
      )
  
      # Validate restore pass
      transformer_containers_r = []
      for i in range(n_layers):
          transformer_containers_r.append(
              TransformerLayerContainer(model)
          )
  
      non_transformer_container_r = NonTransformerContainer(model)
  
      restore_inference_model(
          buffer,
          metadata,
          transformer_containers_r,
          non_transformer_container_r,
      )
  
      validate_containers(
          transformer_containers_r,
          non_transformer_container_r,
          transformer_params,
          non_transformer_params,
      )
  
file_path: DeepSpeed/tests/unit/inference/v2/model_implementations/parameters/test_contiguify.py
imports: from typing import List
  import pytest
  import torch
  from deepspeed.accelerator import get_accelerator
  from deepspeed.inference.v2.model_implementations.flat_model_helpers import flatten_inference_model, restore_inference_model
  from deepspeed.inference.v2.model_implementations.layer_container_base import LayerContainer
  from .utils import SimpleParam, DummyInferenceModel
label_desc: Use image node_id 1 for calling a global function with example usage: test_contiguify_roundtrip() without return types