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
  from typing import Optional
  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 HfFolder, Repository, create_repo, whoami
  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