code: def run_local_whitening_test(
      vec_size,
      num_rows,
      chunk_size,
      window_size,
      partition_size,
      initial_num_partitions,
  ):
      rng = default_rng()
      data = rng.normal(size=(num_rows, vec_size))
      mt = hl.utils.range_matrix_table(
          num_rows, vec_size, n_partitions=initial_num_partitions
      )
      mt = mt.annotate_globals(data=data)
      mt = mt.annotate_entries(x=mt.data[mt.row_idx, mt.col_idx])
  
      tsm = _make_tsm(
          mt.x,
          chunk_size,
          partition_size=partition_size,
          whiten_window_size=window_size,
      )
      ht = tsm.block_table
      whitened_hail = np.vstack(
          ht.aggregate(hl.agg.collect(tsm.block_expr))
      )
      whitened_naive = naive_whiten(data.T, window_size)
      np.testing.assert_allclose(
          whitened_hail, whitened_naive, rtol=1e-04
      )
  
file_path: hail/hail/python/test/hail/experimental/test_local_whitening.py
imports: import numpy
  from numpy.random import default_rng
  from hail.methods.pca import _make_tsm
  import hail
  from ..helpers import *
label_desc: Use image node_id 2 for calling a global function with example usage: run_local_whitening_test(vec_size, num_rows, chunk_size, window_size, partition_size, initial_num_partitions) without return types