for diags, (vecs, banded_mat) in tests.items():
          mask = banded_mat[0] == 0
          input_mat = np.random.randint(10, size=mask.shape)
          solution = input_mat * mask + banded_mat[0]
          output = array_ops.matrix_set_diag(