code: def _matrix_diag_part_v2(
      input, k, padding_value, name=None
  ):  # pylint: disable=redefined-builtin,unused-argument
      """Implements tf.raw_ops.MatrixDiagPartV2, for scalar k."""
      if np.array(k).ndim > 0:
          raise NotImplementedError
      shp = np.shape(input)
  
      if JAX_MODE:
          if len(shp) > 2:
              from jax import (
                  vmap,
              )  # pylint: disable=g-import-not-at-top
  
              return vmap(_matrix_diag_part_v2, (0, None, None))(
                  input, k, padding_value
              )
          return np.diag(input, k=k)
  
      input = np.reshape(input, (-1, shp[-2], shp[-1]))
      output = np.array([np.diag(arr, k=k) for arr in input])
      return output.reshape(*(shp[:-2] + output.shape[1:]))
  
imports: import collections
  import numpy
  import numpy
  from tensorflow_probability.python.internal.backend.numpy import _utils
label_desc: Use image node_id 2 for calling a global function with example usage: _matrix_diag_part_v2(input, k, padding_value, name) and returns: output, np, vmap