TF Attention

attention.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow.keras.layers import Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "attn = tf.keras.layers.Attention(use_scale=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = tf.random.uniform((1, 2, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "def call(q, v, k, mask_q=None, mask_v=None):\n",
    "    \"\"\" Call attention instance \"\"\"\n",
    "    attn = tf.keras.layers.Attention(use_scale=True)\n",
    "    return attn(inputs=[q, v, k], mask=[mask_q, mask_v])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: id=123, shape=(1, 2, 2), dtype=float32, numpy=\n",
       "array([[[0.62968266, 0.6612503 ],\n",
       "        [0.6235384 , 0.73767066]]], dtype=float32)>"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "call(x, x, x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: id=135, shape=(1, 2, 2), dtype=float32, numpy=\n",
       "array([[[0.62968266, 0.6612503 ],\n",
       "        [0.6235384 , 0.73767066]]], dtype=float32)>"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "call(q=x, v=x, k=x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyAttention(tf.keras.Model):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super(MyAttention, self).__init__()\n",
    "        self.attention = Attention(use_scale=True)\n",
    "        \n",
    "    def call(self, q, v, k, mask_q=None, mask_v=None):\n",
    "        return self.attention(inputs=[q, v, k], mask=[mask_q, mask_v])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "needs_attention = NeedsAttention()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: id=147, shape=(1, 2, 2), dtype=float32, numpy=\n",
       "array([[[0.62968266, 0.6612503 ],\n",
       "        [0.6235384 , 0.73767066]]], dtype=float32)>"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "needs_attention(x, x, x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "__call__() missing 1 required positional argument: 'inputs'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-33-5fa3b47998d9>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mneeds_attention\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mq\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mv\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mk\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m: __call__() missing 1 required positional argument: 'inputs'"
     ]
    }
   ],
   "source": [
    "needs_attention(q=x, v=x, k=x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: id=157, shape=(1, 2, 2), dtype=float32, numpy=\n",
       "array([[[0.62968266, 0.6612503 ],\n",
       "        [0.6235384 , 0.73767066]]], dtype=float32)>"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "needs_attention.call(q=x, v=x, k=x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}