code: def ref_matrix_diag(
      x,
  ):
      y = np.zeros(x.shape + (x.shape[-1],))
      if x.ndim == 2:
          for b in range(x.shape[0]):
              y[b, :, :] = np.diag(x[b, :])
      if x.ndim == 3:
          for t in range(x.shape[0]):
              for b in range(x.shape[1]):
                  y[t, b, :, :] = np.diag(x[t, b, :])
      return y
  
imports: import pytest
  import numpy
  import nnabla.functions
  from nbla_test_utils import list_context
label_desc: Use image node_id 1 for calling a global function with example usage: ref_matrix_diag(x) and returns: y