code: def generate_skewed_distribution(alphabet_size, total_counts, seed):
      rng = np.random.default_rng(seed)
  
      base_counts = []
      for _ in range(alphabet_size):
          base_counts.append(rng.integers(low=1, high=alphabet_size))
  
      base_counts = np.array(base_counts)
      arr_counts = (
          base_counts / np.sum(base_counts) * total_counts
      ).astype(np.int32)
      count_diff = total_counts - np.sum(arr_counts)
  
      if count_diff > 0:
          arr_counts[0] += count_diff
      elif count_diff < 0:
          for ind in range(len(arr_counts)):
              if arr_counts[ind] > -1 * count_diff + 1:
                  arr_counts[ind] += count_diff
  
      return arr_counts
  
file_path: NeuralCompression/projects/jax_entropy_coders/test_entropy_coders.py
imports: import craystack
  import jax.numpy
  import jax_arithmetic_coder
  import numpy
  import pytest
  import scipy
label_desc: Use image node_id 4 for calling a global function with example usage: generate_skewed_distribution(alphabet_size, total_counts, seed) and returns: arr_counts