def Init(shape, rng):
    """Returns orthogonalized random normal values with the given `shape`."""
    # Have at least 2 elements in shape.
    cur_shape = list(shape)
    while len(cur_shape) < 2:
      cur_shape = [1] + cur_shape

    # Flatten the input shape with the last dimension remaining.
    n_rows = 1
    for dim in cur_shape[:-1]:
      n_rows *= dim
    n_cols = cur_shape[-1]
    flat_shape = (n_cols, n_rows) if n_rows < n_cols else (n_rows, n_cols)

    # Generate a random matrix
    a = random.normal(rng, flat_shape, dtype=jnp.float32)

    # Compute the qr factorization
    q, r = jnp.linalg.qr(a)

    # Make Q uniform
    d = jnp.diag(r)
    q *= jnp.sign(d)

    # Transpose and reshape back q if needed.
    if n_rows < n_cols:
      q = jnp.transpose(q)
    q = jnp.reshape(q, shape)

    # Return scaled as requested.
    return stddev * q
Returns orthogonalized random normal values with the given `shape`.
trax/layers/initializers.py