code: def _tensor_diag(
      diagonal, name=None
  ):  # pylint: disable=redefined-outer-name
      """tf.linalg.tensor_diag reimpl."""
      del name
      if diagonal.ndim == 0:
          return diagonal
      out = np.zeros(
          (diagonal.shape[0], diagonal.shape[0])
          + diagonal.shape[1:]
          + diagonal.shape[1:],
          dtype=diagonal.dtype,
      )
      for i in range(diagonal.shape[0]):
          if JAX_MODE:
              out = out.at[i, i].set(_tensor_diag(diagonal[i]))
          else:
              out[i, i] = _tensor_diag(diagonal[i])
      axes = (
          [0]
          + list(range(2, diagonal.ndim + 1))
          + [1]
          + list(range(diagonal.ndim + 1, diagonal.ndim * 2))
      )
      return np.transpose(out, axes)
  
imports: import collections
  import functools
  import numpy
  from tensorflow_probability.python.internal.backend.numpy import _utils
  from tensorflow_probability.python.internal.backend.numpy import ops
label_desc: Use image node_id 22 for calling a global function with example usage: _tensor_diag(diagonal, name) and returns: np, diagonal