code: def pmin(x: jax.Array, axis_name: str = "device"):
      """Compute a multi-device reduce min on x over the device axis axis_name."""
      return lax.pmin(x, axis_name)
  
imports: from functools import partial
  import jax
  from jax import lax
label_desc: Use image node_id 3 for calling a global function with example usage: pmin(x, axis_name) and returns: lax