code: def get_sqrt_alpha_prod(
      state: CommonSchedulerState,
      original_samples: jnp.ndarray,
      noise: jnp.ndarray,
      timesteps: jnp.ndarray,
  ):
      alphas_cumprod = state.alphas_cumprod
  
      sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
      sqrt_alpha_prod = sqrt_alpha_prod.flatten()
      sqrt_alpha_prod = broadcast_to_shape_from_left(
          sqrt_alpha_prod, original_samples.shape
      )
  
      sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
      sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
      sqrt_one_minus_alpha_prod = broadcast_to_shape_from_left(
          sqrt_one_minus_alpha_prod, original_samples.shape
      )
  
      return sqrt_alpha_prod, sqrt_one_minus_alpha_prod
  
imports: import importlib
  import math
  import os
  from dataclasses import dataclass
  from enum import Enum
  from typing import Optional, Tuple, Union
  import flax
  import jax.numpy
  from ..utils import BaseOutput, PushToHubMixin
label_desc: Use image node_id 3 for calling a global function with example usage: get_sqrt_alpha_prod(state, original_samples, noise, timesteps) and returns: sqrt_alpha_prod,  sqrt_one_minus_alpha_prod