code: def sample_cxcywha(key, batch_shape):
      """Sample rotated bounding boxes [cx, cy, w, h, a (radians)]."""
      scale = jnp.array([0.3, 0.3, 0.5, 0.5, 1.0])
      offset = jnp.array([0.35, 0.35, 0, 0, 0])
      return (
          jax.random.uniform(key, shape=(*batch_shape, 5)) * scale
          + offset
      )
  
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  import jax.numpy
  import numpy
  from scenic.model_lib.base_models import box_utils
  from shapely import geometry
label_desc: Use image node_id 2 for calling a global function with example usage: sample_cxcywha(key, batch_shape) and returns: unknown