(rng, n_rows=256, n_columns=0, scale_by_norm=True,
                   normalize_data=False, epsilon=0.0001)
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.
Generator for approximate softmax orthogonal kernel feature matrix.
Generator for approximate softmax orthogonal kernel feature matrix .
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)
def get_2d_array ( rng , n_rows = 256 , n_columns = 0 , scale_by_norm = True , normalize_data = False , epsilon = 0.0001 ) : 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 )