code: def Gax(x, ax):
      
  
      slc = (slice(None),)*ax + (slice(-1, None),)
      xg = np.roll(x, -1, axis=ax) - x
      xg[slc] = 0.0
      return xg
docs: Compute gradient of `x` along axis `ax`.
  
      Parameters
      ----------
      x : array_like
        Input array
      ax : int
        Axis on which gradient is to be computed
  
      Returns
      -------
      xg : ndarray
        Output array