code: def interpolate_pos_embedding(
      state_dict,
      out_seqlen,
      pos_embedding_name="model.pos_encoder.pe",
      interleave=False,
  ):
      orig_emb = state_dict["state_dict"][pos_embedding_name]
      assert (
          out_seqlen % orig_emb.shape[1]
      ) == 0, "out_seqlen must be a multiple of the original sequence length"
      reps = [1 for i in orig_emb.shape]
      reps[1] = out_seqlen // orig_emb.shape[1]
  
      if interleave:
          assert (
              math.isqrt(orig_emb.shape[1]) ** 2 == orig_emb.shape[1]
          ), "interleave only works for square lengths"
          assert (
              math.isqrt(out_seqlen) ** 2 == out_seqlen
          ), "interleave only works for square lengths"
          assert (
              math.isqrt(reps[1]) ** 2 == reps[1]
          ), "out_seqlen / seqlen must be a perfect square"
  
          emb_square = rearrange(
              orig_emb,
              "b (h w) d -> b h w d",
              h=math.isqrt(orig_emb.shape[1]),
          )
          emb_square_expanded = emb_square.repeat_interleave(
              math.isqrt(reps[1]), axis=1
          ).repeat_interleave(math.isqrt(reps[1]), axis=2)
          new_emb = rearrange(
              emb_square_expanded, "b h w d -> b (h w) d"
          )
          state_dict["state_dict"][pos_embedding_name] = new_emb
      else:
          state_dict["state_dict"][
              pos_embedding_name
          ] = orig_emb.repeat(*reps)
  
      ret = remove_model_prefix(state_dict)
      # # HACK: this is a hack for block-sparse flash attention
      ret = {
          k: v
          for k, v in ret.items()
          if not k.endswith("inner_attn.layout")
      }
      return ret
  
file_path: flash-attention/training/src/utils/checkpoint.py
imports: import re
  from pathlib import Path
  import torch
  import math
  from einops import rearrange
label_desc: Use image node_id 3 for calling a global function with example usage: interpolate_pos_embedding(state_dict, out_seqlen, pos_embedding_name, interleave) and returns: ret