index.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Embedding(6, 3, padding_idx=5)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([[0., 1., 1.],\n",
       "        [2., 3., 3.],\n",
       "        [4., 5., 5.],\n",
       "        [6., 7., 7.],\n",
       "        [8., 9., 9.],\n",
       "        [0., 0., 0.]], requires_grad=True)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "tensor([[  0.,   1.,   0.,  10., 100.,   0.],\n",
       "        [  1.,   0.,   0.,  10.,   0.,   0.]])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "A = torch.Tensor([\n",
    "    [0, 1, 0, 10, 100, 0],\n",
    "    [1, 0, 0, 10, 0, 0]\n",
    "])\n",
    "\n",
    "embeddings = torch.Tensor([\n",
    "    [0, 1, 1,],\n",
    "    [2, 3, 3],\n",
    "    [4, 5, 5],\n",
    "    [6, 7, 7],\n",
    "    [8 ,9, 9],\n",
    "])\n",
    "\n",
    "def add_pad_embedding(embeddings):\n",
    "    return torch.cat([embeddings, torch.zeros(1, embeddings.shape[1])])\n",
    "\n",
    "embeddings = add_pad_embedding(embeddings)\n",
    "embeddings = torch.nn.Embedding.from_pretrained(embeddings, freeze=False, padding_idx=-1)\n",
    "display(embeddings)\n",
    "display(embeddings.weight)\n",
    "display(A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "1. turn nonzeros into indices\n",
    "2. use indices as in normal transformer to get embedding,\n",
    "pretty much like token indices - should I include padding tokens?\n",
    "3. but also somehow use these indices to extract the output\n",
    "4. it may be that I can also accomplish this via masking\n",
    "'''\n",
    "def pad(vec, seq_len, pad_token=-1):\n",
    "    return torch.cat([\n",
    "        vec, torch.Tensor([pad_token]).expand(seq_len - vec.size()[0])\n",
    "    ])\n",
    "\n",
    "def select_nonzero(A, seq_len=10, pad_token=-1):\n",
    "    seqs = []\n",
    "    for row in A:\n",
    "        nonzeros = pad(row.nonzero().T[0], seq_len, pad_token)\n",
    "        seqs.append(nonzeros)\n",
    "    return torch.stack(seqs)\n",
    "\n",
    "def select_nonzero_vec(seq, seq_len=10, pad_token=-1):\n",
    "    return seq.gather(1, (seq == 0).sort(dim=1, stable=True)[1])\n",
    "\n",
    "def index_tensor(seq, indices):\n",
    "    return seq.T[indices.long()].permute(0,2,1).diagonal().T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[  0.,   1.,   0.,  10., 100.,   0.],\n",
       "        [  1.,   0.,   0.,  10.,   0.,   0.]])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "tensor([[0, 1],\n",
       "        [0, 3],\n",
       "        [0, 4],\n",
       "        [1, 0],\n",
       "        [1, 3]])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "display(A)\n",
    "A.nonzero()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[  0.,   1.,   0.,  10., 100.,   0.],\n",
       "        [  1.,   0.,   0.,  10.,   0.,   0.]])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "tensor([[ 1.,  3.,  4., -1., -1., -1., -1., -1., -1., -1.],\n",
       "        [ 0.,  3., -1., -1., -1., -1., -1., -1., -1., -1.]])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "tensor([[  1.,  10., 100.,   0.,   0.,   0.],\n",
       "        [  1.,  10.,   0.,   0.,   0.,   0.]])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "nonzeros = select_nonzero(A)\n",
    "nonzeros_vec = select_nonzero_vec(A)\n",
    "display(A)\n",
    "display(nonzeros)\n",
    "display(nonzeros_vec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[  0.,   1.,   0.,  10., 100.,   0.],\n",
       "        [  1.,   0.,   0.,  10.,   0.,   0.]])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 1.,  3.,  4., -1., -1., -1., -1., -1., -1., -1.],\n",
       "        [ 0.,  3., -1., -1., -1., -1., -1., -1., -1., -1.]])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nonzeros"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "index_attempt = A.T[nonzeros.long()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[  1.,  10., 100.,   0.,   0.,   0.,   0.,   0.,   0.,   0.],\n",
       "        [  1.,  10.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.]])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "index_attempt.permute(0,2,1).diagonal().T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[  1.,  10., 100.,   0.,   0.,   0.,   0.,   0.,   0.,   0.],\n",
       "        [  1.,  10.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.]])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "index_tensor(A, nonzeros)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 3., 4., 5., 5., 5., 5., 5., 5., 5.],\n",
       "        [0., 3., 5., 5., 5., 5., 5., 5., 5., 5.]])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nonzeros = select_nonzero(A, pad_token=embeddings.weight.size()[0]-1)\n",
    "nonzeros"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[2., 3., 3.],\n",
       "         [6., 7., 7.],\n",
       "         [8., 9., 9.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.]],\n",
       "\n",
       "        [[0., 1., 1.],\n",
       "         [6., 7., 7.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.]]], grad_fn=<EmbeddingBackward0>)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "squished_embeddings = embeddings(nonzeros.long())\n",
    "squished_embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 10, 3])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "squished_embeddings.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[2., 3., 3.],\n",
       "         [6., 7., 7.],\n",
       "         [8., 9., 9.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.]],\n",
       "\n",
       "        [[0., 1., 1.],\n",
       "         [6., 7., 7.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.],\n",
       "         [0., 0., 0.]]], grad_fn=<EmbeddingBackward0>)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "tensor([[1., 3., 4., 5., 5., 5., 5., 5., 5., 5.],\n",
       "        [0., 3., 5., 5., 5., 5., 5., 5., 5., 5.]])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(squished_embeddings)\n",
    "display(nonzeros)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[  2.,   3.,   3.],\n",
       "         [ 60.,  70.,  70.],\n",
       "         [800., 900., 900.],\n",
       "         [  0.,   0.,   0.],\n",
       "         [  0.,   0.,   0.],\n",
       "         [  0.,   0.,   0.],\n",
       "         [  0.,   0.,   0.],\n",
       "         [  0.,   0.,   0.],\n",
       "         [  0.,   0.,   0.],\n",
       "         [  0.,   0.,   0.]],\n",
       "\n",
       "        [[  0.,   1.,   1.],\n",
       "         [ 60.,  70.,  70.],\n",
       "         [  0.,   0.,   0.],\n",
       "         [  0.,   0.,   0.],\n",
       "         [  0.,   0.,   0.],\n",
       "         [  0.,   0.,   0.],\n",
       "         [  0.,   0.,   0.],\n",
       "         [  0.,   0.,   0.],\n",
       "         [  0.,   0.,   0.],\n",
       "         [  0.,   0.,   0.]]], grad_fn=<MulBackward0>)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "squished_embeddings * index_tensor(A, nonzeros).unsqueeze(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 2.,  3.,  3.],\n",
       "         [18., 21., 21.],\n",
       "         [32., 36., 36.],\n",
       "         [ 0.,  0.,  0.],\n",
       "         [ 0.,  0.,  0.],\n",
       "         [ 0.,  0.,  0.],\n",
       "         [ 0.,  0.,  0.],\n",
       "         [ 0.,  0.,  0.],\n",
       "         [ 0.,  0.,  0.],\n",
       "         [ 0.,  0.,  0.]],\n",
       "\n",
       "        [[ 0.,  0.,  0.],\n",
       "         [18., 21., 21.],\n",
       "         [ 0.,  0.,  0.],\n",
       "         [ 0.,  0.,  0.],\n",
       "         [ 0.,  0.,  0.],\n",
       "         [ 0.,  0.,  0.],\n",
       "         [ 0.,  0.,  0.],\n",
       "         [ 0.,  0.,  0.],\n",
       "         [ 0.,  0.,  0.],\n",
       "         [ 0.,  0.,  0.]]], grad_fn=<MulBackward0>)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "squished_embeddings * nonzeros.unsqueeze(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 10, 3])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "squished_embeddings.shape"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "decrypt",
   "language": "python",
   "name": "decrypt"
  },
  "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}