test_np_pytorch_storage_views.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test that we can make copy-free views on an underlying storage and copy-free assignments to those views - for numpy and pytorch views using numpy storage\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastai.text import * "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#numpy: verify that we can make copy-free views on an underlying storage and copy-free assignments to those views\n",
    "bs,seq_len  = 32,10\n",
    "nToks       = bs*(seq_len+1)\n",
    "np_storage  = np.zeros(bs*2*(seq_len+1), dtype=np.int64)\n",
    "\n",
    "#create views on the storage\n",
    "npbatchview = np_storage[:nToks].reshape(bs,-1)\n",
    "x,y         = npbatchview[:,0:seq_len], npbatchview[:,1:seq_len+1]\n",
    "\n",
    "np.add(np_storage, 1, out=np_storage)\n",
    "\n",
    "#assert that the modification of the storage is seen by the views\n",
    "assert x.flatten().all() and y.flatten().all() \n",
    "\n",
    "#assert that a change to a slice in x outside the scope of y does not affect y\n",
    "x[:,0] = 2\n",
    "assert (x[:,0]==2).all() and (y == 1).all()   \n",
    "\n",
    "#assert that a change to a slice in y outside the scope of x does not affect x\n",
    "y[:,seq_len-1] = 3\n",
    "assert (x[:,0]==2).all() and \\\n",
    "       (x[:,1:]==1).all() and \\\n",
    "       (y[:,seq_len-1] == 3).all() and \\\n",
    "       (y[:,seq_len-2] == 1).all()\n",
    "\n",
    "#assert that x and y only modifed the first part of the underlying storage\n",
    "np_storage_view = np_storage.reshape(2*bs,-1)\n",
    "assert (np_storage_view[bs:,:]==1).all() \n",
    "\n",
    "#assert that modifcation to x and y appears in the underlying storage at the expected positions\n",
    "assert (np_storage_view[:bs,0]        ==2).all() and \\\n",
    "       (np_storage_view[:bs,1:seq_len]==1).all() and \\\n",
    "       (np_storage_view[:bs,seq_len]  ==3).all() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#xt,yt      = torch.as_tensor(batchview[:,0:seq_len], dtype=torch.long), torch.as_tensor(batchview[:,1:seq_len+1], dtype=torch.long)\n",
    "#batchview  = torch.as_tensor( np_storage, dtype=torch.long ).reshape(bs,-1)\n",
    "#print(xt.storage())\n",
    "#print(yt.storage())\n",
    "bs,seq_len  = 32,10\n",
    "nToks       = bs*(seq_len+1)\n",
    "#int64 because pytorch embedding-look-up requires it. pytorch see this storage as torch.LongStorage\n",
    "np_storage  = np.zeros(bs*2*(seq_len+1), dtype=np.int64)  \n",
    "npbatchview = np_storage[:nToks].reshape(bs,-1)\n",
    "np_x,np_y   = npbatchview[:,0:seq_len], npbatchview[:,1:seq_len+1]\n",
    "\n",
    "#the following pt_storage uses np_storage as storage so it s is actually a view. It increases the refrencecount\n",
    "#We could set np_storage=None AFTER the following statement and pt_storage would still hold on to the allocated storage\n",
    "pt_storage  = torch.from_numpy( np_storage ) \n",
    "\n",
    "\n",
    "ptbatchview = pt_storage[:nToks].reshape(bs,-1)\n",
    "pt_x,pt_y   = ptbatchview[:,0:seq_len], ptbatchview[:,1:seq_len+1]\n",
    "assert type(ptbatchview) == torch.Tensor and type(pt_x) == torch.Tensor and type(pt_y) == torch.Tensor \n",
    "\n",
    "np.add(np_storage, 1, out=np_storage)\n",
    "\n",
    "#assert that the modification of the storage is seen by the views\n",
    "assert (pt_x.flatten()==1).all() and (pt_y.flatten()==1).all() \n",
    "\n",
    "#assert that a change to a slice in pt_x outside the scope of pt_y does not affect pt_y\n",
    "pt_x[:,0] = 2\n",
    "assert (pt_x[:,0]==2).all() and (pt_y == 1).all()   \n",
    "\n",
    "#assert that a change to a slice in pt_y outside the scope of pt_x does not affect pt_x\n",
    "pt_y[:,seq_len-1] = 3\n",
    "assert (pt_x[:,0]==2).all() and \\\n",
    "       (pt_x[:,1:]==1).all() and \\\n",
    "       (pt_y[:,seq_len-1] == 3).all() and \\\n",
    "       (pt_y[:,seq_len-2] == 1).all()\n",
    "\n",
    "#assert that a change to a slice in pt_y outside the scope of pt_x does not affect pt_x\n",
    "pt_y[:,seq_len-1] = 3\n",
    "assert (pt_x[:,0]==2).all() and \\\n",
    "       (pt_x[:,1:]==1).all() and \\\n",
    "       (pt_y[:,seq_len-1] == 3).all() and \\\n",
    "       (pt_y[:,seq_len-2] == 1).all()\n",
    "\n",
    "#assert that x and y only modifed the first part of the underlying storage\n",
    "pt_storage_view = pt_storage.reshape(2*bs,-1)\n",
    "assert (pt_storage_view[bs:,:]==1).all() \n",
    "\n",
    "#assert that modifcation to x and y appears in the underlying pt_storage at the expected positions\n",
    "assert (pt_storage_view[:bs,0]        ==2).all() and \\\n",
    "       (pt_storage_view[:bs,1:seq_len]==1).all() and \\\n",
    "       (pt_storage_view[:bs,seq_len]  ==3).all() \n",
    "\n",
    "#assert that x and y only modifed the first part of the underlying storage\n",
    "np_storage_view = np_storage.reshape(2*bs,-1)\n",
    "assert (np_storage_view[bs:,:]==1).all() \n",
    "\n",
    "#assert that modifcation to x and y appears in the underlying storage at the expected positions\n",
    "assert (np_storage_view[:bs,0]        ==2).all() and \\\n",
    "       (np_storage_view[:bs,1:seq_len]==1).all() and \\\n",
    "       (np_storage_view[:bs,seq_len]  ==3).all() \n",
    "\n",
    "#assert that creating a tensor from np_x and np_y gives the same view as x and y\n",
    "pt_x2,pt_y2 = torch.from_numpy(np_x), torch.from_numpy(np_y)\n",
    "np_x[0,0]   = 99\n",
    "np_y[-1,-1] = 999\n",
    "assert pt_x[0,0] == 99 and pt_y[-1,-1]==999 and pt_x2[0,0] == 99 and pt_y2[-1,-1]==999"
   ]
  }
 ],
 "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}