Vision transformer

import sys import tensorflow as tf from tensorflow.keras import Model, layers, initializers import numpy as np from tensorflow import keras save_attn_