code: def _set_diag(
      input, diagonal, name=None, k=0, align="RIGHT_LEFT"
  ):  # pylint: disable=unused-argument,redefined-builtin
      if k != 0 or align != "RIGHT_LEFT":
          raise NotImplementedError(
              "set_diag not implemented for k != 0"
          )
      return np.where(
          np.eye(diagonal.shape[-1], dtype=np.bool_),
          diagonal[..., np.newaxis, :],
          input,
      )
  
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 18 for calling a global function with example usage: _set_diag(input, diagonal, name, k, align) and returns: np