code: def _diag(a, xp):
      s0, s1 = a.shape
      ret = xp.zeros((s0, s1, s1), a.dtype)
      arange_s1 = numpy.arange(s1)
      ret[:, arange_s1, arange_s1] = a
      return ret
  
imports: import numpy
  from chainer import backend
  from chainer import function_node
  from chainer.utils import argument
  from chainer.utils import type_check
label_desc: Use image node_id 3 for calling a global function with example usage: _diag(a, xp) and returns: ret