code: def aot_compile(
      prompt=default_prompt,
      negative_prompt=default_neg_prompt,
      seed=default_seed,
      guidance_scale=default_guidance_scale,
      num_inference_steps=default_num_steps,
  ):
      prompt_ids, neg_prompt_ids = tokenize_prompt(
          prompt, negative_prompt
      )
      prompt_ids, neg_prompt_ids, rng = replicate_all(
          prompt_ids, neg_prompt_ids, seed
      )
      g = jnp.array(
          [guidance_scale] * prompt_ids.shape[0], dtype=jnp.float32
      )
      g = g[:, None]
  
      return (
          pmap(
              pipeline._generate,
              static_broadcasted_argnums=[3, 4, 5, 9],
          )
          .lower(
              prompt_ids,
              p_params,
              rng,
              num_inference_steps,  # num_inference_steps
              height,  # height
              width,  # width
              g,
              None,
              neg_prompt_ids,
              False,  # return_latents
          )
          .compile()
      )
  
file_path: llmga/llmga/diffusers/examples/research_projects/sdxl_flax/sdxl_single_aot.py
imports: import time
  import jax
  import jax.numpy
  import numpy
  from flax.jax_utils import replicate
  from jax import pmap
  from jax.experimental.compilation_cache import compilation_cache
  from diffusers import FlaxStableDiffusionXLPipeline
label_desc: Use image node_id 3 for calling a global function with example usage: aot_compile(prompt, negative_prompt, seed, guidance_scale, num_inference_steps) and returns: pmap