code: def create_directory_if_not_exists(path_str: str) -> str:
      """Creates a directory if it doesn't exist, and returns the directory path."""
      if os.path.isdir(path_str):
          return path_str
      elif (
          input(
              f"{path_str} does not exist, create directory? [y/n]"
          ).lower()
          == "y"
      ):
          os.makedirs(path_str)
          return path_str
      else:
          raise NotADirectoryError(path_str)
  
file_path: serve/examples/large_models/inferentia2/util/inf2_save_split_checkpoints.py
imports: import argparse
  import os
  import torch
  from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
  from transformers.models.opt import OPTForCausalLM
  from transformers_neuronx.module import save_pretrained_split
label_desc: Use image node_id 1 for calling a global function with example usage: create_directory_if_not_exists(path_str) and returns: path_str, path_str