code: def test_diag_indices():
      di = diag_indices(4)
      a = np.array(
          [
              [1, 2, 3, 4],
              [5, 6, 7, 8],
              [9, 10, 11, 12],
              [13, 14, 15, 16],
          ]
      )
      a[di] = 100
      assert_array_equal(
          a,
          np.array(
              [
                  [100, 2, 3, 4],
                  [5, 100, 7, 8],
                  [9, 10, 100, 12],
                  [13, 14, 15, 100],
              ]
          ),
      )
  
      # Now, we create indices to manipulate a 3-d array:
      d3 = diag_indices(2, 3)
  
      # And use it to set the diagonal of a zeros array to 1:
      a = np.zeros((2, 2, 2), int)
      a[d3] = 1
      assert_array_equal(
          a, np.array([[[1, 0], [0, 0]], [[0, 0], [0, 1]]])
      )
  
imports: import pytest
  import numpy
  from numpy.testing import assert_, assert_equal, assert_array_equal, assert_almost_equal, assert_array_almost_equal, assert_raises, assert_raises_regex
  from numpy.lib.index_tricks import mgrid, ogrid, ndenumerate, fill_diagonal, diag_indices, diag_indices_from, index_exp, ndindex, r_, s_, ix_
label_desc: Use image node_id 2 for calling a global function with example usage: test_diag_indices() without return types