{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Operations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Most basic operations and torch operations work on proxies and are added to the computation graph.\n",
    "\n",
    "In this example we get the sum of the hidden states and add them to the hidden_states themselves (for whatever reason). By saving the various steps, we can see how the values change."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You're using a GPT2TokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n"
     ]
    }
   ],
   "source": [
    "from nnsight import LanguageModel\n",
    "import torch \n",
    "\n",
    "model = LanguageModel('openai-community/gpt2', device_map='cuda')\n",
    "\n",
    "with model.trace('The Eiffel Tower is in the city of') as tracer:\n",
    "\n",
    "    hidden_states_pre = model.transformer.h[-1].output[0].save()\n",
    "\n",
    "    hs_sum = torch.sum(hidden_states_pre).save()\n",
    "\n",
    "    hs_edited = hidden_states_pre + hs_sum\n",
    "\n",
    "    hs_edited = hs_edited.save()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[ 0.0505, -0.1728, -0.1690,  ..., -1.0096,  0.1280, -1.0687],\n",
      "         [ 8.7494,  2.9057,  5.3024,  ..., -8.0418,  1.2964, -2.8677],\n",
      "         [ 0.2960,  4.6686, -3.6642,  ...,  0.2391, -2.6064,  3.2263],\n",
      "         ...,\n",
      "         [ 2.1537,  6.8917,  3.8651,  ...,  0.0588, -1.9866,  5.9188],\n",
      "         [-0.4460,  7.4285, -9.3065,  ...,  2.0528, -2.7946,  0.5556],\n",
      "         [ 6.6286,  1.7258,  4.7969,  ...,  7.6714,  3.0683,  2.0481]]],\n",
      "       device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(501.2959, device='cuda:0', grad_fn=<SumBackward0>)\n",
      "tensor([[[501.3464, 501.1231, 501.1269,  ..., 500.2863, 501.4239, 500.2272],\n",
      "         [510.0453, 504.2016, 506.5983,  ..., 493.2541, 502.5923, 498.4282],\n",
      "         [501.5919, 505.9645, 497.6317,  ..., 501.5350, 498.6895, 504.5222],\n",
      "         ...,\n",
      "         [503.4496, 508.1876, 505.1610,  ..., 501.3547, 499.3093, 507.2147],\n",
      "         [500.8499, 508.7244, 491.9894,  ..., 503.3487, 498.5013, 501.8515],\n",
      "         [507.9245, 503.0217, 506.0928,  ..., 508.9673, 504.3641, 503.3440]]],\n",
      "       device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    }
   ],
   "source": [
    "print(hidden_states_pre)\n",
    "print(hs_sum)\n",
    "print(hs_edited)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ndif",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
