code: def pad_to_expected_shape(
      x: torch.Tensor, expected_sharded_ten: ShardedTensor
  ):
      pad_args = []
      assert len(x.shape) == len(expected_sharded_ten.local_shape)
      # Reversed iteration order because F.pad expects so
      for x_sh, exp_sh, axis_fragm in reversed(
          list(
              zip(
                  x.shape,
                  expected_sharded_ten.local_shape,
                  expected_sharded_ten.axis_fragmentations,
              )
          )
      ):
          if x_sh == exp_sh:
              pad_args.extend((0, 0))
          elif x_sh > exp_sh:
              assert (
                  False
              ), f"Expected shape ({exp_sh}) smaller than actual ({x_sh}) for {repr(expected_sharded_ten)}"
          else:
              pad_args.extend((0, exp_sh - x_sh))
      # TODO: behavior control with envvar is for testing purposes only, remove it
      if not int(os.environ.get("DIST_CKPT_PAD_REPLICATE", 0)):
          return torch.nn.functional.pad(x, pad_args)
  
      # unsqueeze and squeeze to get shapes supported by cudnn
      print(f"Replicating last row for {expected_sharded_ten.key}")
      if x.dtype == torch.bfloat16:
          return (
              torch.nn.functional.pad(
                  x.float().unsqueeze(0), pad_args, mode="replicate"
              )
              .squeeze(0)
              .bfloat16()
          )
      return torch.nn.functional.pad(
          x.unsqueeze(0), pad_args, mode="replicate"
      ).squeeze(0)
  
imports: import logging
  import os
  from functools import partial
  from logging import getLogger
  from pathlib import Path
  from typing import Callable, List, Optional, Tuple
  import numpy
  import torch
  import zarr
  from ..core import CheckpointingException
  from ..dict_utils import dict_list_map_inplace
  from ..mapping import ShardedStateDict, ShardedTensor, is_main_replica
  from .base import LoadShardedStrategy, SaveShardedStrategy, StrategyAction, default_strategies
label_desc: Use image node_id 9 for calling a global function with example usage: pad_to_expected_shape(x, expected_sharded_ten) and returns: torch, torch, torch