lucidrains

import torch
import torch.nn as nn
import torch.nn.functional as F

# helpers

def make_unit_length(x, epsilon=1e-6):
    norm = x.norm(p=2, dim=-1, keepdim=True)
    return x.div(norm + epsilon)