code: def _extract_into_tensor(arr, timesteps, broadcast_shape):
      """
      Extract values from a 1-D numpy array for a batch of indices.
  
      :param arr: the 1-D numpy array.
      :param timesteps: a tensor of indices into the array to extract.
      :param broadcast_shape: a larger shape of K dimensions with the batch
                              dimension equal to the length of timesteps.
      :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
      """
      if not isinstance(arr, torch.Tensor):
          arr = torch.from_numpy(arr)
      res = arr[timesteps].float().to(timesteps.device)
      while len(res.shape) < len(broadcast_shape):
          res = res[..., None]
      return res.expand(broadcast_shape)
  
imports: import argparse
  import inspect
  import logging
  import math
  import os
  from pathlib import Path
  import accelerate
  import datasets
  import torch
  import torch.nn.functional
  from accelerate import Accelerator
  from accelerate.logging import get_logger
  from accelerate.utils import ProjectConfiguration
  from datasets import load_dataset
  from huggingface_hub import create_repo, upload_folder
  from onnxruntime.training.optim.fp16_optimizer import FP16_Optimizer
  from onnxruntime.training.ortmodule import ORTModule
  from packaging import version
  from torchvision import transforms
  from tqdm.auto import tqdm
  import diffusers
  from diffusers import DDPMPipeline, DDPMScheduler, UNet2DModel
  from diffusers.optimization import get_scheduler
  from diffusers.training_utils import EMAModel
  from diffusers.utils import check_min_version, is_accelerate_version, is_tensorboard_available, is_wandb_available
  from diffusers.utils.import_utils import is_xformers_available
label_desc: Use image node_id 1 for calling a global function with example usage: _extract_into_tensor(arr, timesteps, broadcast_shape) and returns: res