yuanzhi-zhu

"""
This code was originally obtained from:
https://github.com/meta-llama/codellama/blob/main/llama/model.py
adapted from https://github.com/naver-ai/rope-vit/blob/main/self-attn/rope_self_attn.py
"""

import torch
import torch.nn as nn
from functools import partial
import einops