code: def sample_cxcywh_bbox(key, batch_shape):
      """Samples a bounding box in the [cx, cy, w, h] in [0, 1] range format."""
      frac = 0.8
      sample = jax.random.uniform(key, shape=(*batch_shape, 4)) * frac
      cx, cy, w, h = jnp.split(sample, indices_or_sections=4, axis=-1)
      # Make sure the bounding box doesn't cross the right and top image borders
      w = jnp.where(cx + w / 2.0 >= 1.0, frac * 2.0 * (1.0 - cx), w)
      h = jnp.where(cy + h / 2.0 >= 1.0, frac * 2.0 * (1.0 - cy), h)
      # Make sure the bounding box doesn't cross the left and bottom image borders
      w = jnp.where(cx - w / 2.0 <= 0.0, frac * 2.0 * cx, w)
      h = jnp.where(cy - h / 2.0 <= 0.0, frac * 2.0 * cy, h)
  
      bbox = jnp.concatenate([cx, cy, w, h], axis=-1)
      return bbox
  
file_path: scenic/scenic/projects/baselines/detr/tests/test_train_utils.py
imports: import collections
  import os
  from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  import jax.numpy
  import numpy
  from scenic.projects.baselines.detr import train_utils
label_desc: Use image node_id 1 for calling a global function with example usage: sample_cxcywh_bbox(key, batch_shape) and returns: bbox