code: def _zero_pad(x, pad, axis):  # pylint: disable = invalid-name
      """Helper for jnp.pad with 0s for single-axis case."""
      pad_widths = [(0, 0)] * len(x.shape)
      pad_widths[axis] = pad  # Padding on axis.
      return jnp.pad(x, pad_widths, mode="constant")
  
imports: import functools
  import itertools
  import operator
  from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import initializers
label_desc: Use image node_id 2 for calling a global function with example usage: _zero_pad(x, pad, axis) and returns: jnp