Function sqrt

SwissArmyTransformer/sat/mpu/utils.py:81–82

(x)

            per_partition_vocab_size, rank, world_size)

def sqrt(x):
    return int(math.sqrt(x) + 1e-4)

def unscaled_init_method(sigma):
    """Init method based on N(0, sigma)."""