code: def repack_diagonals_in_tests(tests, align=None):
      # The original test cases are LEFT_LEFT aligned.
      if align == default_v2_alignment or align is None:
          return tests
  
      new_tests = dict()
      # Loops through each case.
      for diag_index, (
          packed_diagonals,
          padded_diagonals,
      ) in tests.items():
          num_rows, num_cols = padded_diagonals.shape[-2:]
          repacked_diagonals = repack_diagonals(
              packed_diagonals,
              diag_index,
              num_rows,
              num_cols,
              align=align,
          )
          new_tests[diag_index] = (repacked_diagonals, padded_diagonals)
  
      return new_tests
  
imports: import numpy
  from tensorflow.compiler.tests import xla_test
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import gen_array_ops
  from tensorflow.python.platform import googletest
label_desc: Use image node_id 3 for calling a global function with example usage: repack_diagonals_in_tests(tests, align) and returns: new_tests, tests