examples/research_projects/performer/modeling_flax_performer_utils.py/GaussianOrthogonalRandomMatrix/get_2d_array
class GaussianOrthogonalRandomMatrix:    def get_2d_array(self):
        nb_full_blocks = int(self.nb_rows / self.nb_columns)
        block_list = []
        rng = self.key
        for _ in range(nb_full_blocks):
            rng, rng_input = jax.random.split(rng)
            unstructured_block = random.normal(rng_input, (self.nb_columns, self.nb_columns))
            q, _ = jnp.linalg.qr(unstructured_block)
            q = jnp.transpose(q)
            block_list.append(q)
        remaining_rows = self.nb_rows - nb_full_blocks * self.nb_columns
        if remaining_rows > 0:
            rng, rng_input = jax.random.split(rng)
            unstructured_block = random.normal(rng_input, (self.nb_columns, self.nb_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 self.scaling == 0:
            multiplier = jnp.linalg.norm(random.normal(self.key, (self.nb_rows, self.nb_columns)), axis=1)
        elif self.scaling == 1:
            multiplier = jnp.sqrt(float(self.nb_columns)) * jnp.ones((self.nb_rows))
        else:
            raise ValueError("Scaling must be one of {0, 1}. Was %s" % self._scaling)

        return jnp.matmul(jnp.diag(multiplier), final_matrix)