Skip to content

Instantly share code, notes, and snippets.

"""Quantized multi head attention layer."""
import math
import tensorflow as tf
from google3.third_party.tensorflow.python.ops import array_ops
from google3.third_party.tensorflow.python.ops import math_ops
from google3.third_party.tensorflow.python.ops import special_math_ops
from google3.third_party.tensorflow_models.official.projects.qat.nlp.quantization import helper