def get_2d_array(rng, n_rows=256, n_columns=0, scale_by_norm=True,
                   normalize_data=False, epsilon=0.0001):
    """Generator for approximate softmax orthogonal kernel feature matrix.

    Args:
      rng: Random number generator.
      n_rows: Number of rows.
      n_columns: Number of columns.
      scale_by_norm: Boolean; whether to scale orthogonal random matrix.
      normalize_data: predicate indicating whether data should be normalized.
      epsilon: numerical stabilizer.

    Returns:
      Orthogonal kernel feature matrix.
    """
    n_full_blocks = int(n_rows / n_columns)
    block_list = []
    rng_key = rng
    for _ in range(n_full_blocks):
      rng, rng_input = random.split(rng)
      unstructured_block = random.normal(rng_input, (n_columns, n_columns))
      q, _ = jnp.linalg.qr(unstructured_block)
      q = jnp.transpose(q)
      block_list.append(q)
    remaining_rows = n_rows - n_full_blocks * n_columns
    if remaining_rows > 0:
      rng, rng_input = random.split(rng)
      unstructured_block = random.normal(rng_input, (n_columns, n_columns))
      q, _ = jnp.linalg.qr(unstructured_block)
      q = jnp.transpose(q)
      block_list.append(q[0:remaining_rows])
    final_matrix = jnp.vstack(block_list)

    if scale_by_norm:
      multiplier = jnp.linalg.norm(
          random.normal(rng_key, (n_rows, n_columns)), axis=1)
    else:
      multiplier = jnp.sqrt(float(n_columns)) * jnp.ones((n_rows))

    return jnp.matmul(jnp.diag(multiplier), final_matrix)
Generator for approximate softmax orthogonal kernel feature matrix.

    Args:
      rng: Random number generator.
      n_rows: Number of rows.
      n_columns: Number of columns.
      scale_by_norm: Boolean; whether to scale orthogonal random matrix.
      normalize_data: predicate indicating whether data should be normalized.
      epsilon: numerical stabilizer.

    Returns:
      Orthogonal kernel feature matrix.
    
trax/layers/research/sparsity.py