code: def _zero_pad(x, pad, axis):
      """Helper for np.pad with 0s for single-axis case."""
      pad_widths = [(0, 0)] * len(x.shape)
      pad_widths[axis] = pad  # Padding on axis.
      return np.pad(
          x,
          pad_widths,
          mode="constant",
          constant_values=x.dtype.type(0),
      )
  
imports: import collections
  import os
  import gin
  import gym
  import numpy
  from trax import fastmath
  from trax.rl import advantages
  from trax.supervised import training
label_desc: Use image node_id 2 for calling a global function with example usage: _zero_pad(x, pad, axis) and returns: np