code: def set_diag(x, new_diag):
      """Set the diagonal along the last two axis.
  
      Parameters
      ----------
      x : array-like, shape=[dim]
          Initial array.
      new_diag : array-like, shape=[dim[-2]]
          Values to set on the diagonal.
  
      Returns
      -------
      None
  
      Notes
      -----
      This mimics tensorflow.linalg.set_diag(x, new_diag), when new_diag is a
      1-D array, but modifies x instead of creating a copy.
      """
      arr_shape = x.shape
      x[..., range(arr_shape[-2]), range(arr_shape[-1])] = new_diag
      return x
  
imports: from ._dispatch import BACKEND_NAME, _common
  from ._dispatch import numpy
label_desc: Use image node_id 15 for calling a global function with example usage: set_diag(x, new_diag) and returns: x