code: def save_sharded_to_hf_format(
      input_ckpt: Union[FluidPath, str],
      params: dict,
      output_path: Union[FluidPath, str],
      cpu: bool = False,
      dtype: str = "fp16",
  ):
      devices = np.array([jax.devices()[0]]).reshape((1, 1))
      with maps.mesh(devices, ("dp", "mp")):
          params_local = params.copy()
          params_local[
              "cores_per_replica"
          ] = maps.thread_resources.env.shape["mp"]
          network = CausalTransformer(params_local)
  
          save_pytree_as_hf(
              network.state,
              input_ckpt=input_ckpt,
              shards_in=params["cores_per_replica"],
              output_path=output_path,
              n_layers=params["layers"],
              np_dtype=np_dtype,
              torch_dtype=torch_dtype,
              n_seq=params["seq"],
          )
  
imports: import argparse
  import io
  import multiprocessing
  import time
  import warnings
  import os
  import re
  from typing import Iterable, List, Union
  import json
  import jax
  import jax.numpy
  from jax.experimental import maps
  from pathy import FluidPath, Pathy
  import numpy
  import optax
  import torch
  from tqdm import tqdm
  from mesh_transformer.transformer_shard import CausalTransformer
label_desc: Use image node_id 13 for calling a global function with example usage: save_sharded_to_hf_format(input_ckpt, params, output_path, cpu, dtype) without return types