code: def diag_einsum(
      input_subscripts, output_subscript, *ioperands, **kwargs
  ):
      (output_shape,) = utils.argument.parse_kwargs(
          kwargs, ("output_shape", None)
      )
      return einsum.DiagEinSum(
          in_subs=input_subscripts,
          out_sub=output_subscript,
          out_shape=output_shape,
      ).apply(ioperands)[0]
  
imports: import unittest
  import numpy
  import chainer
  from chainer.backends import cuda
  from chainer.functions.math import einsum
  from chainer import gradient_check
  from chainer import testing
  from chainer.testing import attr
  from chainer import utils
label_desc: Use image node_id 4 for calling a global function with example usage: diag_einsum(input_subscripts, output_subscript) and returns: einsum