code: def save_exported_model(
      layout,
      model,
      model_H: Hyperparams,
      save_dir,
      save_heads: Set[str],
  ):
      """
      Exporting a model allows it to be run with a different layout than it was trained with.
  
      Currently, uploading/loading an exported model is slower than saving/restoring a checkpoint,
      but if we can get exporting to be sufficiently fast, then we could replace legacy_checkpoints.py with
      this "exporting" approach.
      """
      if blobs.is_blob_url(save_dir):
          local_dir = os.path.join("/tmp", str(uuid.uuid4()))
      else:
          local_dir = save_dir
      os.makedirs(os.path.join(local_dir, "checkpoint"), exist_ok=True)
  
      def export_fine_piece(
          fine_model_piece_dict: dict, chkpt_prefix: str
      ):
          fine_piece_path = os.path.join(
              local_dir,
              "checkpoint",
              f"{chkpt_prefix}_shard_{layout.shard_idx:03d}.pkl",
          )
          # print(f"Uploading fine_piece: {fine_piece_path}")
          torch.save(fine_model_piece_dict, fine_piece_path)
          torch.cuda.synchronize()  # Verify that the piece has finished being written
  
      # Export the embeddings
      if model.include_input_embeddings:
          export_fine_piece(
              model.embedding.state_dict(), "input_embeddings"
          )
  
          if model.include_pos_embeddings:
              export_fine_piece(
                  model.position_embedding.state_dict(),
                  "position_embedding",
              )
      # Export the resblocks
      for resblock_idx, resblock in enumerate(model.torso.resblocks):
          export_fine_piece(
              resblock.state_dict(), f"resblock_{resblock_idx:04d}"
          )
      # Export the final_layer_norm
      if model.include_final_layer_norm:
          export_fine_piece(model.ln_f.state_dict(), "final_layer_norm")
      # Export the unembeddings
      if model.include_output_unembeddings:
          export_fine_piece(
              {"unembedding_weights": model.unembedding_weights},
              "output_unembeddings",
          )
      for head in save_heads:
          export_fine_piece(
              model.scalar_heads[head].state_dict(),
              f"output_head_{head}",
          )
      if blobs.is_blob_url(save_dir):
          blobs.parallel_copy_recursive(local_dir, save_dir)
          shutil.rmtree(local_dir)
  
imports: import json
  import os
  import re
  import shutil
  import uuid
  from dataclasses import dataclass, field
  from typing import Callable, Optional, Set, List, Dict
  import blobfile
  import numpy
  import torch
  from mpi4py import MPI
  import summarize_from_feedback
  from summarize_from_feedback.model_layout import ModelLayout
  from summarize_from_feedback.models import sample_fns
  from summarize_from_feedback.models.loss_functions import softmax_xent_loss_fn
  from summarize_from_feedback.models.transformer import Hyperparams
  from summarize_from_feedback.models.transformer import build_with_random_weights
  from summarize_from_feedback.utils import exact_div, hyperparams
  from summarize_from_feedback.utils import blobs
  from summarize_from_feedback.utils.dist_utils import setup_cuda_device_and_dist, create_data_parallel_comm, create_within_replica_comm
  from summarize_from_feedback.utils.nested import map_nested
  from summarize_from_feedback.utils.torch_utils import nans, to_numpy
label_desc: Use image node_id 2 for calling a global function with example usage: save_exported_model(layout, model, model_H, save_dir, save_heads) without return types