{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-06-27T13:36:50.722271Z","iopub.execute_input":"2021-06-27T13:36:50.722954Z","iopub.status.idle":"2021-06-27T13:37:16.216457Z","shell.execute_reply.started":"2021-06-27T13:36:50.722811Z","shell.execute_reply":"2021-06-27T13:37:16.215357Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#!conda install -c conda-forge pillow -y\n#!conda install -c conda-forge pydicom -y\n#!conda install -c conda-forge gdcm -y\n#!pip install pylibjpeg pylibjpeg-libjpeg\n#!pip install pydicom","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install --upgrade numpy==1.20.0","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%bash\nconda install -c conda-forge gdcm -y","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install pylibjpeg pylibjpeg-libjpeg","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import cv2\nimport io\nimport matplotlib\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport os\nimport time\nimport argparse\nimport pandas as pd\nfrom PIL import Image\nimport pydicom\nimport torch.nn as nn\nimport seaborn as sns\nimport warnings","metadata":{"execution":{"iopub.status.busy":"2021-06-27T13:47:13.120796Z","iopub.execute_input":"2021-06-27T13:47:13.121208Z","iopub.status.idle":"2021-06-27T13:47:15.010807Z","shell.execute_reply.started":"2021-06-27T13:47:13.121174Z","shell.execute_reply":"2021-06-27T13:47:15.009832Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"TRAIN_DIR = '../input/siim-covid19-detection/train'\nTEST_DIR = '../input/siim-covid19-detection/test'","metadata":{"execution":{"iopub.status.busy":"2021-06-27T13:47:18.854924Z","iopub.execute_input":"2021-06-27T13:47:18.855430Z","iopub.status.idle":"2021-06-27T13:47:18.860505Z","shell.execute_reply.started":"2021-06-27T13:47:18.855385Z","shell.execute_reply":"2021-06-27T13:47:18.858666Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train = pd.read_csv('../input/siim-covid19-detection/train_image_level.csv')\ntrain_study = pd.read_csv('../input/siim-covid19-detection/train_study_level.csv')\n#train_study.head()","metadata":{"execution":{"iopub.status.busy":"2021-06-27T13:47:21.591425Z","iopub.execute_input":"2021-06-27T13:47:21.592064Z","iopub.status.idle":"2021-06-27T13:47:21.640932Z","shell.execute_reply.started":"2021-06-27T13:47:21.592022Z","shell.execute_reply":"2021-06-27T13:47:21.640154Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_study['id'] = train_study['id'].apply(lambda i: i.split('_')[0])\ntrain_study.rename(columns={'Negative for Pneumonia': '0','Typical Appearance': '1',\"Indeterminate Appearance\": '2',\n                   \"Atypical Appearance\": \"3\"}, inplace=True)\n#train_study.head()","metadata":{"execution":{"iopub.status.busy":"2021-06-27T13:47:24.227455Z","iopub.execute_input":"2021-06-27T13:47:24.227858Z","iopub.status.idle":"2021-06-27T13:47:24.239153Z","shell.execute_reply.started":"2021-06-27T13:47:24.227821Z","shell.execute_reply":"2021-06-27T13:47:24.237698Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"labels = []\ndef get_label(row):\n    for c in train_study.columns:\n        if row[c] == 1:\n            labels.append(int(c))\n            \ntrain_study.apply(get_label, axis=1)\ntrain_study.drop(columns=['0', '1','2', '3'], inplace=True)\ntrain_study['label'] = labels\n#train_study.head()","metadata":{"execution":{"iopub.status.busy":"2021-06-27T13:47:32.140303Z","iopub.execute_input":"2021-06-27T13:47:32.140670Z","iopub.status.idle":"2021-06-27T13:47:32.382573Z","shell.execute_reply.started":"2021-06-27T13:47:32.140639Z","shell.execute_reply":"2021-06-27T13:47:32.381502Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from os import listdir, walk\nfrom skimage import exposure\nimport torch\nfrom torch.utils.data import Dataset, DataLoader","metadata":{"execution":{"iopub.status.busy":"2021-06-27T13:47:36.140181Z","iopub.execute_input":"2021-06-27T13:47:36.140750Z","iopub.status.idle":"2021-06-27T13:47:36.341094Z","shell.execute_reply.started":"2021-06-27T13:47:36.140701Z","shell.execute_reply":"2021-06-27T13:47:36.340011Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CovLungDataset(Dataset):\n    def __init__(self, dir_path, labels_data, transforms=None, new_size=(512, 512)):\n        self.dir_path = dir_path\n        self.labels_data = labels_data\n        self.new_size = new_size\n        self.transforms = transforms\n\n    def __len__(self):\n        return len(self.labels_data)\n\n    def __getitem__(self, idx):\n        if torch.is_tensor(idx):\n            idx = idx.tolist()\n        \n        image_dir = self.labels_data.iloc[idx]['id']\n        img_label = self.labels_data.iloc[idx]['label']\n        \n        path_to_img = os.path.join(self.dir_path, image_dir)\n        # get first image path only\n        path_to_img = os.path.join(path_to_img, listdir(path_to_img)[0])\n        path_to_img = os.path.join(path_to_img, next(walk(path_to_img))[2][0])\n        \n        # read image\n        data = pydicom.dcmread(path_to_img)\n        image = data.pixel_array\n        image = exposure.equalize_hist(image)\n        \n        good_height, good_width = self.new_size\n        image = cv2.resize(image, (good_width, good_height), interpolation=Image.LANCZOS)\n        \n        # data augmentation\n        if self.transforms:\n            # doesn't work on floats\n            image = (image * 255).astype(np.uint8)\n            image = self.transforms(image=image)['image']\n        \n        sample = {'image': image, 'label': img_label}\n        return sample\n","metadata":{"execution":{"iopub.status.busy":"2021-06-27T13:47:38.958760Z","iopub.execute_input":"2021-06-27T13:47:38.959136Z","iopub.status.idle":"2021-06-27T13:47:38.970404Z","shell.execute_reply.started":"2021-06-27T13:47:38.959104Z","shell.execute_reply":"2021-06-27T13:47:38.969194Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\ntransform = A.Compose([A.RandomBrightnessContrast(brightness_limit=[-0.2, 0.2], \n                                                  contrast_limit=[-0.2, 0.2], \n                                                  p=0.5),\n                       A.ShiftScaleRotate(scale_limit=[-0.1, 0.3], \n                                          shift_limit=0.1, \n                                          rotate_limit=20, \n                                          border_mode=cv2.BORDER_CONSTANT,\n                                          p=0.5),\n                       # reshape image of size (k, n, 1) into (1, k, n)\n                       ToTensorV2(p=1.0)\n                      ])\n","metadata":{"execution":{"iopub.status.busy":"2021-06-27T13:47:44.462756Z","iopub.execute_input":"2021-06-27T13:47:44.463160Z","iopub.status.idle":"2021-06-27T13:47:45.355837Z","shell.execute_reply.started":"2021-06-27T13:47:44.463114Z","shell.execute_reply":"2021-06-27T13:47:45.354766Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\n\ntransformed_train_dataset = CovLungDataset(dir_path=TRAIN_DIR,\n                                     labels_data=train_study[['id', 'label']],\n                                     transforms=transform,\n                                     new_size=(512, 512))\n\n\ntrain_dataloader = DataLoader(transformed_train_dataset, batch_size=8, shuffle=False, num_workers=2)\n\n","metadata":{"execution":{"iopub.status.busy":"2021-06-27T13:47:48.769652Z","iopub.execute_input":"2021-06-27T13:47:48.770041Z","iopub.status.idle":"2021-06-27T13:47:48.777745Z","shell.execute_reply.started":"2021-06-27T13:47:48.770010Z","shell.execute_reply":"2021-06-27T13:47:48.776503Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\ntransformed_test_dataset = CovLungDataset(dir_path=TEST_DIR,\n                                     labels_data=train_study[['id', 'label']],\n                                     transforms=transform,\n                                     new_size=(512, 512))\n\n\ntest_dataloader = DataLoader(transformed_test_dataset, batch_size=8, shuffle=False, num_workers=2)\n","metadata":{"execution":{"iopub.status.busy":"2021-06-27T13:47:53.245804Z","iopub.execute_input":"2021-06-27T13:47:53.246179Z","iopub.status.idle":"2021-06-27T13:47:53.254586Z","shell.execute_reply.started":"2021-06-27T13:47:53.246145Z","shell.execute_reply":"2021-06-27T13:47:53.253211Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nfrom mpl_toolkits.axes_grid1 import ImageGrid\nimport numpy as np\n\n\ndataiter = iter(train_dataloader)\nsample = dataiter.next()\n#sample = next(iter(transformed_train_dataset))\nimages = sample['image'] \nlabels = sample['label']\n\nfig = plt.figure(figsize=(15., 25.))\ngrid = ImageGrid(fig, 111, \n                 nrows_ncols=(2, 4),\n                 axes_pad=.4,\n                 )\n\nlabels_dict = {0: 'Negative for Pneumonia',  1: 'Typical Appearance',  2: 'Indeterminate Appearance',  3: 'Atypical Appearance'}\n\nj = 0\nfor ax, im in zip(grid, images):\n    im = im.numpy()\n    im = np.transpose(im, (1,2,0))\n    ax.imshow(im, cmap='gray')\n    ax.set_title(labels_dict[labels[j].item()], fontsize=12)\n    j += 1\n\nplt.show()\n","metadata":{"execution":{"iopub.status.busy":"2021-06-27T13:47:56.162312Z","iopub.execute_input":"2021-06-27T13:47:56.162722Z","iopub.status.idle":"2021-06-27T13:48:05.418640Z","shell.execute_reply.started":"2021-06-27T13:47:56.162681Z","shell.execute_reply":"2021-06-27T13:48:05.417425Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"parser = argparse.ArgumentParser()\nparser.add_argument('--num_epochs', type=int, default=10, help='Number of training epochs')\nparser.add_argument('--batch_size', type=int, default=64, help='Batch size')\nparser.add_argument('--lr', type=float, default=5e-4, help='Learning rate')\nparser.add_argument('--l2_reg', type=float, default=0, help='L2 regularisation')\nparser.add_argument('--aug', action='store_true', default=False, help='Use data augmentation')\nparser.add_argument('--data_path', type=str, default='/input/siim-covid19-detection',help='Path to data.')\nparser.add_argument('--bond_dim', type=int, default=5, help='MPS Bond dimension')\nparser.add_argument('--nChannel', type=int, default=1, help='Number of input channels')\nparser.add_argument('--dense_net', action='store_true', default=False, help='Using Dense Net model')","metadata":{"execution":{"iopub.status.busy":"2021-06-27T13:48:15.816083Z","iopub.execute_input":"2021-06-27T13:48:15.816462Z","iopub.status.idle":"2021-06-27T13:48:15.830945Z","shell.execute_reply.started":"2021-06-27T13:48:15.816431Z","shell.execute_reply":"2021-06-27T13:48:15.829666Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"args = parser.parse_args([])","metadata":{"execution":{"iopub.status.busy":"2021-06-27T13:48:21.008626Z","iopub.execute_input":"2021-06-27T13:48:21.009005Z","iopub.status.idle":"2021-06-27T13:48:21.015011Z","shell.execute_reply.started":"2021-06-27T13:48:21.008972Z","shell.execute_reply":"2021-06-27T13:48:21.013171Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Miscellaneous initialization\ntorch.manual_seed(0)\nstart_time = time.time()\n\n# MPS parameters\nbond_dim = 20\nadaptive_mode = False\nperiodic_bc = False\n\n# Training parameters\n#num_train = 2000\n#num_test = 1000\nbatch_size = 100\nnum_epochs = args.num_epochs\nlearn_rate = args.lr\nl2_reg = args.l2_reg\n\n\nbatch_size = args.batch_size\nbond_dim = args.bond_dim\n\n\n# LoTeNet parameters\nadaptive_mode = False \nperiodic_bc   = False\n\nkernel = 2 # Stride along spatial dimensions\noutput_dim = 4 # output dimension\n \nfeature_dim = 2\n\n#logFile = time.strftime(\"%Y%m%d_%H_%M\")+'.txt'\n#makeLogFile(logFile)\n\nnormTensor = 0.5*torch.ones(args.nChannel)\n### Data processing and loading....","metadata":{"execution":{"iopub.status.busy":"2021-06-27T14:20:53.576740Z","iopub.execute_input":"2021-06-27T14:20:53.577242Z","iopub.status.idle":"2021-06-27T14:20:53.587654Z","shell.execute_reply.started":"2021-06-27T14:20:53.577204Z","shell.execute_reply":"2021-06-27T14:20:53.586306Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def svd_flex(tensor, svd_string, max_D=None, cutoff=1e-10, sv_right=True, sv_vec=None):\n    \"\"\"\n    Split an input tensor into two pieces using a SVD across some partition\n\n    Args:\n        tensor (Tensor):    Pytorch tensor with at least two indices\n\n        svd_string (str):   String of the form 'init_str->left_str,right_str',\n                            where init_str describes the indices of tensor, and\n                            left_str/right_str describe those of the left and\n                            right output tensors. The characters of left_str\n                            and right_str form a partition of the characters in\n                            init_str, but each contain one additional character\n                            representing the new bond which comes from the SVD\n\n                            Reversing the terms in svd_string to the left and\n                            right of '->' gives an ein_string which can be used\n                            to multiply both output tensors to give a (low rank\n                            approximation) of the input tensor\n\n        cutoff (float):     A truncation threshold which eliminates any\n                            singular values which are strictly less than cutoff\n\n        max_D (int):        A maximum allowed value for the new bond. If max_D\n                            is specified, the returned tensors\n\n        sv_right (bool):    The SVD gives two orthogonal matrices and a matrix\n                            of singular values. sv_right=True merges the SV\n                            matrix with the right output, while sv_right=False\n                            merges it with the left output\n\n        sv_vec (Tensor):    Pytorch vector with length max_D, which is modified\n                            in place to return the vector of singular values\n\n    Returns:\n        left_tensor (Tensor),\n        right_tensor (Tensor):  Tensors whose indices are described by the\n                                left_str and right_str parts of svd_string\n\n        bond_dim:               The dimension of the new bond appearing from\n                                the cutoff in our SVD. Note that this generally\n                                won't match the dimension of left_/right_tensor\n                                at this mode, which is padded with zeros\n                                whenever max_D is specified\n    \"\"\"\n\n    def prod(int_list):\n        output = 1\n        for num in int_list:\n            output *= num\n        return output\n\n    with torch.no_grad():\n        # Parse svd_string into init_str, left_str, and right_str\n        svd_string = svd_string.replace(\" \", \"\")\n        init_str, post_str = svd_string.split(\"->\")\n        left_str, right_str = post_str.split(\",\")\n\n        # Check formatting of init_str, left_str, and right_str\n        assert all([c.islower() for c in init_str + left_str + right_str])\n        assert len(set(init_str + left_str + right_str)) == len(init_str) + 1\n        assert len(set(init_str)) + len(set(left_str)) + len(set(right_str)) == len(\n            init_str\n        ) + len(left_str) + len(right_str)\n\n        # Get the special character representing our SVD-truncated bond\n        bond_char = set(left_str).intersection(set(right_str)).pop()\n        left_part = left_str.replace(bond_char, \"\")\n        right_part = right_str.replace(bond_char, \"\")\n\n        # Permute our tensor into something that can be viewed as a matrix\n        ein_str = f\"{init_str}->{left_part+right_part}\"\n        tensor = torch.einsum(ein_str, [tensor]).contiguous()\n\n        left_shape = list(tensor.shape[: len(left_part)])\n        right_shape = list(tensor.shape[len(left_part) :])\n        left_dim, right_dim = prod(left_shape), prod(right_shape)\n\n        tensor = tensor.view([left_dim, right_dim])\n\n        # Get SVD and format so that left_mat * diag(svs) * right_mat = tensor\n        left_mat, svs, right_mat = torch.svd(tensor)\n        svs, _ = torch.sort(svs, descending=True)\n        right_mat = torch.t(right_mat)\n\n        # Decrease or increase our tensor sizes in the presence of max_D\n        if max_D and len(svs) > max_D:\n            svs = svs[:max_D]\n            left_mat = left_mat[:, :max_D]\n            right_mat = right_mat[:max_D]\n        elif max_D and len(svs) < max_D:\n            copy_svs = torch.zeros([max_D])\n            copy_svs[: len(svs)] = svs\n            copy_left = torch.zeros([left_mat.size(0), max_D])\n            copy_left[:, : left_mat.size(1)] = left_mat\n            copy_right = torch.zeros([max_D, right_mat.size(1)])\n            copy_right[: right_mat.size(0)] = right_mat\n            svs, left_mat, right_mat = copy_svs, copy_left, copy_right\n\n        # If given as input, copy singular values into sv_vec\n        if sv_vec is not None and svs.shape == sv_vec.shape:\n            sv_vec[:] = svs\n        elif sv_vec is not None and svs.shape != sv_vec.shape:\n            raise TypeError(\n                f\"sv_vec.shape must be {list(svs.shape)}, but is \"\n                f\"currently {list(sv_vec.shape)}\"\n            )\n\n        # Find the truncation point relative to our singular value cutoff\n        truncation = 0\n        for s in svs:\n            if s < cutoff:\n                break\n            truncation += 1\n        if truncation == 0:\n            raise RuntimeError(\n                \"SVD cutoff too large, attempted to truncate \"\n                \"tensor to bond dimension 0\"\n            )\n\n        # Perform the actual truncation\n        if max_D:\n            svs[truncation:] = 0\n            left_mat[:, truncation:] = 0\n            right_mat[truncation:] = 0\n        else:\n            # If max_D wasn't given, set it to the truncation index\n            max_D = truncation\n            svs = svs[:truncation]\n            left_mat = left_mat[:, :truncation]\n            right_mat = right_mat[:truncation]\n\n        # Merge the singular values into the appropriate matrix\n        if sv_right:\n            right_mat = torch.einsum(\"l,lr->lr\", [svs, right_mat])\n        else:\n            left_mat = torch.einsum(\"lr,r->lr\", [left_mat, svs])\n\n        # Reshape the matrices to make them proper tensors\n        left_tensor = left_mat.view(left_shape + [max_D])\n        right_tensor = right_mat.view([max_D] + right_shape)\n\n        # Finally, permute the indices into the desired order\n        if left_str != left_part + bond_char:\n            left_tensor = torch.einsum(\n                f\"{left_part+bond_char}->{left_str}\", [left_tensor]\n            )\n        if right_str != bond_char + right_part:\n            right_tensor = torch.einsum(\n                f\"{bond_char+right_part}->{right_str}\", [right_tensor]\n            )\n\n        return left_tensor, right_tensor, truncation\n\n\ndef init_tensor(shape, bond_str, init_method):\n    \"\"\"\n    Initialize a tensor with a given shape\n\n    Args:\n        shape:       The shape of our output parameter tensor.\n\n        bond_str:    The bond string describing our output parameter tensor,\n                     which is used in 'random_eye' initialization method.\n                     The characters 'l' and 'r' are used to refer to the\n                     left or right virtual indices of our tensor, and are\n                     both required to be present for the random_eye and\n                     min_random_eye initialization methods.\n\n        init_method: The method used to initialize the entries of our tensor.\n                     This can be either a string, or else a tuple whose first\n                     entry is an initialization method and whose remaining\n                     entries are specific to that method. In each case, std\n                     will always refer to a standard deviation for a random\n                     normal random component of each entry of the tensor.\n\n                     Allowed options are:\n                        * ('random_eye', std): Initialize each tensor input\n                            slice close to the identity\n                        * ('random_zero', std): Initialize each tensor input\n                            slice close to the zero matrix\n                        * ('min_random_eye', std, init_dim): Initialize each\n                            tensor input slice close to a truncated identity\n                            matrix, whose truncation leaves init_dim unit\n                            entries on the diagonal. If init_dim is larger\n                            than either of the bond dimensions, then init_dim\n                            is capped at the smaller bond dimension.\n    \"\"\"\n    # Unpack init_method if it is a tuple\n    if not isinstance(init_method, str):\n        init_str = init_method[0]\n        std = init_method[1]\n        if init_str == \"min_random_eye\":\n            init_dim = init_method[2]\n\n        init_method = init_str\n    else:\n        std = 1e-9\n\n    # Check that bond_str is properly sized and doesn't have repeat indices\n    assert len(shape) == len(bond_str)\n    assert len(set(bond_str)) == len(bond_str)\n\n    if init_method not in [\"random_eye\", \"min_random_eye\", \"random_zero\"]:\n        raise ValueError(f\"Unknown initialization method: {init_method}\")\n\n    if init_method in [\"random_eye\", \"min_random_eye\"]:\n        bond_chars = [\"l\", \"r\"]\n        assert all([c in bond_str for c in bond_chars])\n\n        # Initialize our tensor slices as identity matrices which each fill\n        # some or all of the initially allocated bond space\n        if init_method == \"min_random_eye\":\n\n            # The dimensions for our initial identity matrix. These will each\n            # be init_dim, unless init_dim exceeds one of the bond dimensions\n            bond_dims = [shape[bond_str.index(c)] for c in bond_chars]\n            if all([init_dim <= full_dim for full_dim in bond_dims]):\n                bond_dims = [init_dim, init_dim]\n            else:\n                init_dim = min(bond_dims)\n\n            eye_shape = [init_dim if c in bond_chars else 1 for c in bond_str]\n            expand_shape = [\n                init_dim if c in bond_chars else shape[i]\n                for i, c in enumerate(bond_str)\n            ]\n\n        elif init_method == \"random_eye\":\n            eye_shape = [\n                shape[i] if c in bond_chars else 1 for i, c in enumerate(bond_str)\n            ]\n            expand_shape = shape\n            bond_dims = [shape[bond_str.index(c)] for c in bond_chars]\n\n        eye_tensor = torch.eye(bond_dims[0], bond_dims[1]).view(eye_shape)\n        eye_tensor = eye_tensor.expand(expand_shape)\n\n        tensor = torch.zeros(shape)\n        tensor[[slice(dim) for dim in expand_shape]] = eye_tensor\n\n        # Add on a bit of random noise\n        tensor += std * torch.randn(shape)\n\n    elif init_method == \"random_zero\":\n        tensor = std * torch.randn(shape)\n\n    return tensor\n\n\n### OLDER MISCELLANEOUS FUNCTIONS ###   # noqa: E266\n\n\ndef onehot(labels, max_value):\n    \"\"\"\n    Convert a batch of labels from the set {0, 1,..., num_value-1} into their\n    onehot encoded counterparts\n    \"\"\"\n    label_vecs = torch.zeros([len(labels), max_value])\n\n    for i, label in enumerate(labels):\n        label_vecs[i, label] = 1.0\n\n    return label_vecs\n\n\ndef joint_shuffle(input_data, input_labels):\n    \"\"\"\n    Shuffle input data and labels in a joint manner, so each label points to\n    its corresponding datum. Works for both regular and CUDA tensors\n    \"\"\"\n    assert input_data.is_cuda == input_labels.is_cuda\n    use_gpu = input_data.is_cuda\n    if use_gpu:\n        input_data, input_labels = input_data.cpu(), input_labels.cpu()\n\n    data, labels = input_data.numpy(), input_labels.numpy()\n\n    # Shuffle relative to the same seed\n    np.random.seed(0)\n    np.random.shuffle(data)\n    np.random.seed(0)\n    np.random.shuffle(labels)\n\n    data, labels = torch.from_numpy(data), torch.from_numpy(labels)\n    if use_gpu:\n        data, labels = data.cuda(), labels.cuda()\n\n    return data, labels\n\n\ndef load_HV_data(length):\n    \"\"\"\n    Output a toy \"horizontal/vertical\" data set of black and white\n    images with size length x length. Each image contains a single\n    horizontal or vertical stripe, set against a background\n    of the opposite color. The labels associated with these images\n    are either 0 (horizontal stripe) or 1 (vertical stripe).\n\n    In its current version, this returns two data sets, a training\n    set with 75% of the images and a test set with 25% of the\n    images.\n    \"\"\"\n    num_images = 4 * (2 ** (length - 1) - 1)\n    num_patterns = num_images // 2\n    split = num_images // 4\n\n    if length > 14:\n        print(\n            \"load_HV_data will generate {} images, \"\n            \"this could take a while...\".format(num_images)\n        )\n\n    images = np.empty([num_images, length, length], dtype=np.float32)\n    labels = np.empty(num_images, dtype=np.int)\n\n    # Used to generate the stripe pattern from integer i below\n    template = \"{:0\" + str(length) + \"b}\"\n\n    for i in range(1, num_patterns + 1):\n        pattern = template.format(i)\n        pattern = [int(s) for s in pattern]\n\n        for j, val in enumerate(pattern):\n            # Horizontal stripe pattern\n            images[2 * i - 2, j, :] = val\n            # Vertical stripe pattern\n            images[2 * i - 1, :, j] = val\n\n        labels[2 * i - 2] = 0\n        labels[2 * i - 1] = 1\n\n    # Shuffle and partition into training and test sets\n    np.random.seed(0)\n    np.random.shuffle(images)\n    np.random.seed(0)\n    np.random.shuffle(labels)\n\n    train_images, train_labels = images[split:], labels[split:]\n    test_images, test_labels = images[:split], labels[:split]\n\n    return (\n        torch.from_numpy(train_images),\n        torch.from_numpy(train_labels),\n        torch.from_numpy(test_images),\n        torch.from_numpy(test_labels),\n    )","metadata":{"execution":{"iopub.status.busy":"2021-06-27T13:48:49.050291Z","iopub.execute_input":"2021-06-27T13:48:49.050926Z","iopub.status.idle":"2021-06-27T13:48:49.096535Z","shell.execute_reply.started":"2021-06-27T13:48:49.050890Z","shell.execute_reply":"2021-06-27T13:48:49.095443Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Contractable:\n    \"\"\"\n    Container for tensors with labeled indices and a global batch size\n\n    The labels for our indices give some high-level knowledge of the tensor\n    layout, and permit the contraction of pairs of indices in a more\n    systematic manner. However, much of the actual heavy lifting is done\n    through specific contraction routines in different subclasses\n\n    Attributes:\n        tensor (Tensor):    A Pytorch tensor whose first index is a batch\n                            index. Sub-classes of Contractable may put other\n                            restrictions on tensor\n        bond_str (str):     A string whose letters each label a separate mode\n                            of our tensor, and whose length equals the order\n                            (number of modes) of our tensor\n        global_bs (int):    The batch size associated with all Contractables.\n                            This is shared between all Contractable instances\n                            and allows for automatic expanding of tensors\n    \"\"\"\n\n    # The global batch size\n    global_bs = None\n\n    def __init__(self, tensor, bond_str):\n        shape = list(tensor.shape)\n        num_dim = len(shape)\n        str_len = len(bond_str)\n\n        global_bs = Contractable.global_bs\n        batch_dim = tensor.size(0)\n\n        # Expand along a new batch dimension if needed\n        if (\"b\" not in bond_str and str_len == num_dim) or (\n            \"b\" == bond_str[0] and str_len == num_dim + 1\n        ):\n            if global_bs is not None:\n                tensor = tensor.unsqueeze(0).expand([global_bs] + shape)\n            else:\n                raise RuntimeError(\n                    \"No batch size given and no previous \" \"batch size set\"\n                )\n            if bond_str[0] != \"b\":\n                bond_str = \"b\" + bond_str\n\n        # Check for correct formatting in bond_str\n        elif bond_str[0] != \"b\" or str_len != num_dim:\n            raise ValueError(\n                \"Length of bond string '{bond_str}' \"\n                f\"({len(bond_str)}) must match order of \"\n                f\"tensor ({len(shape)})\"\n            )\n\n        # Set the global batch size if it is unset or needs to be updated\n        elif global_bs is None or global_bs != batch_dim:\n            Contractable.global_bs = batch_dim\n\n        # Check that global batch size agrees with input tensor's first dim\n        elif global_bs != batch_dim:\n            raise RuntimeError(\n                f\"Batch size previously set to {global_bs}\"\n                \", but input tensor has batch size \"\n                f\"{batch_dim}\"\n            )\n\n        # Set the defining attributes of our Contractable\n        self.tensor = tensor\n        self.bond_str = bond_str\n\n    def __mul__(self, contractable, rmul=False):\n        \"\"\"\n        Multiply with another contractable along a linear index\n\n        The default behavior is to multiply the 'r' index of this instance\n        with the 'l' index of contractable, matching the batch ('b')\n        index of both, and take the outer product of other indices.\n        If rmul is True, contractable is instead multiplied on the right.\n        \"\"\"\n        # This method works for general Core subclasses besides Scalar (no 'l'\n        # and 'r' indices), composite contractables (no tensor attribute), and\n        # MatRegion (multiplication isn't just simple index contraction)\n        if (\n            isinstance(contractable, Scalar)\n            or not hasattr(contractable, \"tensor\")\n            or type(contractable) is MatRegion\n        ):\n            return NotImplemented\n\n        tensors = [self.tensor, contractable.tensor]\n        bond_strs = [list(self.bond_str), list(contractable.bond_str)]\n        lowercases = [chr(c) for c in range(ord(\"a\"), ord(\"z\") + 1)]\n\n        # Reverse the order of tensors if needed\n        if rmul:\n            tensors = tensors[::-1]\n            bond_strs = bond_strs[::-1]\n\n        # Check that bond strings are in proper format\n        for i, bs in enumerate(bond_strs):\n            assert bs[0] == \"b\"\n            assert len(set(bs)) == len(bs)\n            assert all([c in lowercases for c in bs])\n            assert (i == 0 and \"r\" in bs) or (i == 1 and \"l\" in bs)\n\n        # Get used and free characters\n        used_chars = set(bond_strs[0]).union(bond_strs[1])\n        free_chars = [c for c in lowercases if c not in used_chars]\n\n        # Rename overlapping indices in the bond strings (except 'b', 'l', 'r')\n        specials = [\"b\", \"l\", \"r\"]\n        for i, c in enumerate(bond_strs[1]):\n            if c in bond_strs[0] and c not in specials:\n                bond_strs[1][i] = free_chars.pop()\n\n        # Combine right bond of left tensor and left bond of right tensor\n        sum_char = free_chars.pop()\n        bond_strs[0][bond_strs[0].index(\"r\")] = sum_char\n        bond_strs[1][bond_strs[1].index(\"l\")] = sum_char\n        specials.append(sum_char)\n\n        # Build bond string of ouput tensor\n        out_str = [\"b\"]\n        for bs in bond_strs:\n            out_str.extend([c for c in bs if c not in specials])\n        out_str.append(\"l\" if \"l\" in bond_strs[0] else \"\")\n        out_str.append(\"r\" if \"r\" in bond_strs[1] else \"\")\n\n        # Build the einsum string for this operation\n        bond_strs = [\"\".join(bs) for bs in bond_strs]\n        out_str = \"\".join(out_str)\n        ein_str = f\"{bond_strs[0]},{bond_strs[1]}->{out_str}\"\n\n        # Contract along the linear dimension to get an output tensor\n        out_tensor = torch.einsum(ein_str, [tensors[0], tensors[1]])\n\n        # Return our output tensor wrapped in an appropriate class\n        if out_str == \"br\":\n            return EdgeVec(out_tensor, is_left_vec=True)\n        elif out_str == \"bl\":\n            return EdgeVec(out_tensor, is_left_vec=False)\n        elif out_str == \"blr\":\n            return SingleMat(out_tensor)\n        elif out_str == \"bolr\":\n            return OutputCore(out_tensor)\n        else:\n            return Contractable(out_tensor, out_str)\n\n    def __rmul__(self, contractable):\n        \"\"\"\n        Multiply with another contractable along a linear index\n        \"\"\"\n        return self.__mul__(contractable, rmul=True)\n\n    def reduce(self):\n        \"\"\"\n        Return the contractable without any modification\n\n        reduce() can be any method which returns a contractable. This is\n        trivially possible for any contractable by returning itself\n        \"\"\"\n        return self\n\n\nclass ContractableList(Contractable):\n    \"\"\"\n    A list of contractables which can all be multiplied together in order\n\n    Calling reduce on a ContractableList instance will first reduce every item\n    to a linear contractable, and then contract everything together\n    \"\"\"\n\n    def __init__(self, contractable_list):\n        # Check that input list is nonempty and has contractables as entries\n        if not isinstance(contractable_list, list) or contractable_list is []:\n            raise ValueError(\"Input to ContractableList must be nonempty list\")\n        for i, item in enumerate(contractable_list):\n            if not isinstance(item, Contractable):\n                raise ValueError(\n                    \"Input items to ContractableList must be \"\n                    f\"Contractable instances, but item {i} is not\"\n                )\n\n        self.contractable_list = contractable_list\n\n    def __mul__(self, contractable, rmul=False):\n        \"\"\"\n        Multiply a contractable by everything in ContractableList in order\n        \"\"\"\n        # The input cannot be a composite contractable\n        assert hasattr(contractable, \"tensor\")\n        output = contractable.tensor\n\n        # Multiply by everything in ContractableList, in the correct order\n        if rmul:\n            for item in self.contractable_list:\n                output = item * output\n        else:\n            for item in self.contractable_list[::-1]:\n                output = output * item\n\n        return output\n\n    def __rmul__(self, contractable):\n        \"\"\"\n        Multiply another contractable by everything in ContractableList\n        \"\"\"\n        return self.__mul__(contractable, rmul=True)\n\n    def reduce(self, parallel_eval=False):\n        \"\"\"\n        Reduce all the contractables in list before multiplying them together\n        \"\"\"\n        c_list = self.contractable_list\n        # For parallel_eval, reduce all contractables in c_list\n        if parallel_eval:\n            c_list = [item.reduce() for item in c_list]\n\n        # Multiply together all the contractables. This multiplies in right to\n        # left order, but certain inefficient contractions are unsupported.\n        # If we encounter an unsupported operation, then try multiplying from\n        # the left end of the list instead\n        while len(c_list) > 1:\n            try:\n                c_list[-2] = c_list[-2] * c_list[-1]\n                del c_list[-1]\n            except TypeError:\n                c_list[1] = c_list[0] * c_list[1]\n                del c_list[0]\n\n        return c_list[0]\n\n\nclass MatRegion(Contractable):\n    \"\"\"\n    A contiguous collection of matrices which are multiplied together\n\n    The input tensor defining our MatRegion must have shape\n    [batch_size, num_mats, D, D], or [num_mats, D, D] when the global batch\n    size is already known\n    \"\"\"\n\n    def __init__(self, mats):\n        shape = list(mats.shape)\n        if len(shape) not in [3, 4] or shape[-2] != shape[-1]:\n            raise ValueError(\n                \"MatRegion tensors must have shape \"\n                \"[batch_size, num_mats, D, D], or [num_mats,\"\n                \" D, D] if batch size has already been set\"\n            )\n\n        super().__init__(mats, bond_str=\"bslr\")\n\n    def __mul__(self, edge_vec, rmul=False):\n        \"\"\"\n        Iteratively multiply an input vector with all matrices in MatRegion\n        \"\"\"\n        # The input must be an instance of EdgeVec\n        if not isinstance(edge_vec, EdgeVec):\n            return NotImplemented\n\n        mats = self.tensor\n        num_mats = mats.size(1)\n\n        # Load our vector and matrix batches\n        dummy_ind = 1 if rmul else 2\n        vec = edge_vec.tensor.unsqueeze(dummy_ind)\n        mat_list = [mat.squeeze(1) for mat in torch.chunk(mats, num_mats, 1)]\n\n        # Do the repeated matrix-vector multiplications in the proper order\n        for i, mat in enumerate(mat_list[:: (1 if rmul else -1)], 1):\n            if rmul:\n                vec = torch.bmm(vec, mat)\n            else:\n                vec = torch.bmm(mat, vec)\n\n        # Since we only have a single vector, wrap it as a EdgeVec\n        return EdgeVec(vec.squeeze(dummy_ind), is_left_vec=rmul)\n\n    def __rmul__(self, edge_vec):\n        return self.__mul__(edge_vec, rmul=True)\n\n    def reduce(self):\n        \"\"\"\n        Multiplies together all matrices and returns resultant SingleMat\n\n        This method uses iterated batch multiplication to evaluate the full\n        matrix product in depth O( log(num_mats) )\n        \"\"\"\n        mats = self.tensor\n        shape = list(mats.shape)\n        size, D = shape[1:3]\n\n        # Iteratively multiply pairs of matrices until there is only one\n        while size > 1:\n            odd_size = size % 2 == 1\n            half_size = size // 2\n            nice_size = 2 * half_size\n\n            even_mats = mats[:, 0:nice_size:2]\n            odd_mats = mats[:, 1:nice_size:2]\n            # For odd sizes, set aside one batch of matrices for the next round\n            leftover = mats[:, nice_size:]\n\n            # Multiply together all pairs of matrices (except leftovers)\n            mats = torch.einsum(\"bslu,bsur->bslr\", [even_mats, odd_mats])\n            mats = torch.cat([mats, leftover], 1)\n\n            size = half_size + int(odd_size)\n\n        # Since we only have a single matrix, wrap it as a SingleMat\n        return SingleMat(mats.squeeze(1))\n\n\nclass OutputCore(Contractable):\n    \"\"\"\n    A single MPS core with a single output index\n    \"\"\"\n\n    def __init__(self, tensor):\n        # Check the input shape\n        if len(tensor.shape) not in [3, 4]:\n            raise ValueError(\n                \"OutputCore tensors must have shape [batch_size, \"\n                \"output_dim, D_l, D_r], or else [output_dim, D_l,\"\n                \" D_r] if batch size has already been set\"\n            )\n\n        super().__init__(tensor, bond_str=\"bolr\")\n\n\nclass SingleMat(Contractable):\n    \"\"\"\n    A batch of matrices associated with a single location in our MPS\n    \"\"\"\n\n    def __init__(self, mat):\n        # Check the input shape\n        if len(mat.shape) not in [2, 3]:\n            raise ValueError(\n                \"SingleMat tensors must have shape [batch_size, \"\n                \"D_l, D_r], or else [D_l, D_r] if batch size \"\n                \"has already been set\"\n            )\n\n        super().__init__(mat, bond_str=\"blr\")\n\n\nclass OutputMat(Contractable):\n    \"\"\"\n    An output core associated with an edge of our MPS\n    \"\"\"\n\n    def __init__(self, mat, is_left_mat):\n        # Check the input shape\n        if len(mat.shape) not in [2, 3]:\n            raise ValueError(\n                \"OutputMat tensors must have shape [batch_size, \"\n                \"D, output_dim], or else [D, output_dim] if \"\n                \"batch size has already been set\"\n            )\n\n        # OutputMats on left edge will have a right-facing bond, and vice versa\n        bond_str = \"b\" + (\"r\" if is_left_mat else \"l\") + \"o\"\n        super().__init__(mat, bond_str=bond_str)\n\n    def __mul__(self, edge_vec, rmul=False):\n        \"\"\"\n        Multiply with an edge vector along the shared linear index\n        \"\"\"\n        if not isinstance(edge_vec, EdgeVec):\n            raise NotImplemented  # noqa: F901\n        else:\n            return super().__mul__(edge_vec, rmul)\n\n    def __rmul__(self, edge_vec):\n        return self.__mul__(edge_vec, rmul=True)\n\n\nclass EdgeVec(Contractable):\n    \"\"\"\n    A batch of vectors associated with an edge of our MPS\n\n    EdgeVec instances are always associated with an edge of an MPS, which\n    requires the is_left_vec flag to be set to True (vector on left edge) or\n    False (vector on right edge)\n    \"\"\"\n\n    def __init__(self, vec, is_left_vec):\n        # Check the input shape\n        if len(vec.shape) not in [1, 2]:\n            raise ValueError(\n                \"EdgeVec tensors must have shape \"\n                \"[batch_size, D], or else [D] if batch size \"\n                \"has already been set\"\n            )\n\n        # EdgeVecs on left edge will have a right-facing bond, and vice versa\n        bond_str = \"b\" + (\"r\" if is_left_vec else \"l\")\n        super().__init__(vec, bond_str=bond_str)\n\n    def __mul__(self, right_vec):\n        \"\"\"\n        Take the inner product of our vector with another vector\n        \"\"\"\n        # The input must be an instance of EdgeVec\n        if not isinstance(right_vec, EdgeVec):\n            return NotImplemented\n\n        left_vec = self.tensor.unsqueeze(1)\n        right_vec = right_vec.tensor.unsqueeze(2)\n        batch_size = left_vec.size(0)\n\n        # Do the batch inner product\n        scalar = torch.bmm(left_vec, right_vec).view([batch_size])\n\n        # Since we only have a single scalar, wrap it as a Scalar\n        return Scalar(scalar)\n\n\nclass Scalar(Contractable):\n    \"\"\"\n    A batch of scalars\n    \"\"\"\n\n    def __init__(self, scalar):\n        # Add dummy dimension if we have a torch scalar\n        shape = list(scalar.shape)\n        if shape is []:\n            scalar = scalar.view([1])\n            shape = [1]\n\n        # Check the input shape\n        if len(shape) != 1:\n            raise ValueError(\n                \"input scalar must be a torch tensor with shape \"\n                \"[batch_size], or [] or [1] if batch size has \"\n                \"been set\"\n            )\n\n        super().__init__(scalar, bond_str=\"b\")\n\n    def __mul__(self, contractable):\n        \"\"\"\n        Multiply a contractable by our scalar and return the result\n        \"\"\"\n        scalar = self.tensor\n        tensor = contractable.tensor\n        bond_str = contractable.bond_str\n\n        ein_string = f\"{bond_str},b->{bond_str}\"\n        out_tensor = torch.einsum(ein_string, [tensor, scalar])\n\n        # Wrap the result in the same class right_contractable belongs to\n        contract_class = type(contractable)\n        if contract_class is not Contractable:\n            return contract_class(out_tensor)\n        else:\n            return Contractable(out_tensor, bond_str)\n\n    def __rmul__(self, contractable):\n        # Scalar multiplication is commutative\n        return self.__mul__(contractable)","metadata":{"execution":{"iopub.status.busy":"2021-06-27T13:48:56.785032Z","iopub.execute_input":"2021-06-27T13:48:56.785421Z","iopub.status.idle":"2021-06-27T13:48:56.843280Z","shell.execute_reply.started":"2021-06-27T13:48:56.785380Z","shell.execute_reply":"2021-06-27T13:48:56.842149Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#from Contractable import (\n#SingleMat,\n#    MatRegion,\n#    OutputCore,\n#    ContractableList,\n#    EdgeVec,\n#    OutputMat,\n#)\n\n\nclass TI_MPS(nn.Module):\n    \"\"\"\n    Sequence MPS which converts input of arbitrary length to a single output vector\n    \"\"\"\n\n    def __init__(\n        self,\n        output_dim,\n        bond_dim,\n        feature_dim=2,\n        parallel_eval=False,\n        fixed_ends=False,\n        init_std=1e-9,\n        use_bias=True,\n        fixed_bias=True,\n    ):\n        super().__init__()\n\n        # Initialize the core tensor defining our model near the identity\n        # This tensor holds all of the trainable parameters of our model\n        tensor = init_tensor(\n            bond_str=\"lri\",\n            shape=[bond_dim, bond_dim, feature_dim],\n            init_method=(\"random_zero\", init_std),\n        )\n        self.register_parameter(name=\"core_tensor\", param=nn.Parameter(tensor))\n\n        # Define our initial vector and terminal matrix, which are both\n        # functional modules, i.e. unchanged during training\n        assert isinstance(fixed_ends, bool)\n        self.init_vector = InitialVector(bond_dim, fixed_vec=fixed_ends)\n        self.terminal_mat = TerminalOutput(bond_dim, output_dim, fixed_mat=fixed_ends)\n\n        # Set the bias matrix\n        if use_bias:\n            # bias_mat is identity when fixed_bias=True, near-identity otherwise\n            if fixed_bias:\n                bias_mat = torch.eye(bond_dim)\n                self.register_buffer(name=\"bias_mat\", tensor=bias_mat)\n            else:\n                bias_mat = init_tensor(\n                    bond_str=\"lr\",\n                    shape=[bond_dim, bond_dim],\n                    init_method=(\"random_eye\", init_std),\n                )\n                self.register_parameter(name=\"bias_mat\", param=nn.Parameter(bias_mat))\n        else:\n            self.bias_mat = None\n\n        # Set the rest of our TI_MPS attributes\n        self.feature_dim = feature_dim\n        self.output_dim = output_dim\n        self.bond_dim = bond_dim\n        self.parallel_eval = parallel_eval\n        self.use_bias = use_bias\n        self.fixed_bias = fixed_bias\n        self.feature_map = None\n\n    def forward(self, input_data):\n        \"\"\"\n        Converts batch input tensor into a batch output tensor\n\n        Args:\n            input_data: A tensor of shape [batch_size, length, feature_dim].\n        \"\"\"\n\n        # Reformat our input to a batch format, padding with zeros as needed\n        batch_input = self.format_input(input_data)\n        batch_size = batch_input.size(0)\n        seq_len = batch_input.size(1)\n\n        # Build up a contractable_list as EdgeVec + MatRegion + OutputMat\n        expanded_core = self.core_tensor.expand(\n            [seq_len, self.bond_dim, self.bond_dim, self.feature_dim]\n        )\n        input_region = InputRegion(\n            expanded_core,\n            use_bias=self.use_bias,\n            fixed_bias=self.fixed_bias,\n            bias_mat=self.bias_mat,\n            ephemeral=True,\n        )\n        contractable_list = [input_region(batch_input)]\n\n        # Prepend an EdgeVec and append an OutputMat\n        contractable_list = [self.init_vector()] + contractable_list\n        contractable_list.append(self.terminal_mat())\n\n        # Wrap contractable_list as a ContractableList instance\n        contractable_list = ContractableList(contractable_list)\n\n        # Contract everything in contractable_list\n        output = contractable_list.reduce(parallel_eval=self.parallel_eval)\n        batch_output = output.tensor\n\n        # Check shape before returning output values\n        assert output.bond_str == \"bo\"\n        assert batch_output.size(0) == batch_size\n        assert batch_output.size(1) == self.output_dim\n\n        return batch_output\n\n    def format_input(self, input_data):\n        \"\"\"\n        Converts input list of sequences into a single batch sequence tensor.\n\n        If input is already a batch tensor, it is returned unchanged. Otherwise,\n        convert input list into a batch sequence with shape [batch_size, length,\n        feature_dim].\n\n        If self.use_bias = self.fixed_bias = True, then sequences of different\n        lengths can be used, in which case shorter sequences are padded with\n        zeros at the end, making the batch tensor length equal to the length\n        of the longest input sequence.\n\n        Args:\n            input_data: A tensor of shape [batch_size, length] or\n            [batch_size, length, feature_dim], or a list of length batch_size,\n            whose i'th item is a tensor of shape [length_i, feature_dim] or\n            [length_i]. If self.use_bias or self.fixed_bias are False, then\n            length_i must be the same for all i.\n        \"\"\"\n        feature_dim = self.feature_dim\n\n        # If we get a batch tensor, just embed it and/or return it unchanged\n        if isinstance(input_data, torch.Tensor):\n            if len(input_data.shape) == 2:\n                input_data = self.embed_input(input_data)\n\n            # Check to make sure shape is alright\n            shape = input_data.shape\n            assert len(shape) == 3\n            assert shape[2] == feature_dim\n\n            return input_data\n\n        # Collate the input list into a single batch tensor\n        elif isinstance(input_data, list):\n            # Check formatting, require that input sequences are either all\n            # unembedded or all pre-embedded\n            num_modes = len(input_data[0].shape)\n            assert num_modes in [1, 2]\n            assert all(\n                [\n                    isinstance(s, torch.Tensor) and len(s.shape) == num_modes\n                    for s in input_data\n                ]\n            )\n            assert num_modes == 1 or all([s.size(1) == feature_dim for s in input_data])\n\n            # Check that all the sequences are the same length or can be padded\n            max_len = max([s.size(0) for s in input_data])\n            can_pad = self.use_bias and self.fixed_bias\n            if not can_pad and any([s.size(0) != max_len for s in input_data]):\n                raise ValueError(\n                    \"To process input_data as list of sequences \"\n                    \"with different lengths, must have self.use_bias=\"\n                    \"self.fixed_bias=True (currently self.use_bias=\"\n                    f\"{self.use_bias}, self.fixed_bias={self.fixed_bias})\"\n                )\n\n            # Pad the sequences with zeros (if needed), return as batch tensor\n            if can_pad:\n                batch_size = len(input_data)\n                full_size = [batch_size, max_len, feature_dim]\n                batch_input = torch.zeros(full_size[: num_modes + 1])\n\n                # Copy each sequence into batch_input\n                for i, seq in enumerate(input_data):\n                    batch_input[i, : seq.size(0)] = seq\n            else:\n                batch_input = torch.stack(input_data)\n\n            # Embed everything (if needed) and return the batch tensor\n            if len(batch_input.shape) == 2:\n                batch_input = self.embed_input(batch_input)\n\n            return batch_input\n\n        else:\n            raise ValueError(\n                \"input_data must either be Tensor with shape\"\n                \"[batch_size, length] or [batch_size, length, feature_dim], \"\n                \"or list of Tensors with shapes [length_i, feature_dim] or \"\n                \"[length_i]\"\n            )\n\n    def embed_input(self, input_data):\n        \"\"\"\n        Embed pixels of input_data into separate local feature spaces\n\n        Args:\n            input_data (Tensor):    Input with shape [batch_size, length].\n\n        Returns:\n            embedded_data (Tensor): Input embedded into a tensor with shape\n                                    [batch_size, input_dim, feature_dim]\n        \"\"\"\n        assert len(input_data.shape) == 2\n\n        # Get relevant dimensions\n        batch_dim, length = input_data.shape\n        feature_dim = self.feature_dim\n        embedded_shape = [batch_dim, length, feature_dim]\n\n        # Apply a custom embedding map if it has been defined by the user\n        if self.feature_map is not None:\n            f_map = self.feature_map\n            embedded_data = torch.stack(\n                [torch.stack([f_map(x) for x in batch]) for batch in input_data]\n            )\n\n            # Make sure our embedded input has the desired size\n            assert list(embedded_data.shape) == embedded_shape\n\n        # Otherwise, use a simple linear embedding map with feature_dim = 2\n        else:\n            if self.feature_dim != 2:\n                raise RuntimeError(\n                    f\"self.feature_dim = {feature_dim}, but \"\n                    \"default feature_map requires self.feature_dim = 2\"\n                )\n            embedded_data = torch.empty(embedded_shape)\n\n            embedded_data[:, :, 0] = input_data\n            embedded_data[:, :, 1] = 1 - input_data\n\n        return embedded_data\n\n    def register_feature_map(self, feature_map):\n        \"\"\"\n        Register a custom feature map to be used for embedding input data\n\n        Args:\n            feature_map (function): Takes a single scalar input datum and\n                                    returns an embedded representation of the\n                                    image. The output size of the function must\n                                    match self.feature_dim. If feature_map=None,\n                                    then the feature map will be reset to a\n                                    simple default linear embedding\n        \"\"\"\n        if feature_map is not None:\n            # Test to make sure feature_map outputs vector of proper size\n            test_out = feature_map(torch.tensor(0))\n            assert isinstance(test_out, torch.Tensor)\n\n            out_shape, needed_shape = list(test_out.shape), [self.feature_dim]\n            if out_shape != needed_shape:\n                raise ValueError(\n                    \"Given feature_map returns values with shape \"\n                    f\"{list(out_shape)}, but should return \"\n                    f\"values of size {list(needed_shape)}\"\n                )\n\n        self.feature_map = feature_map\n\n\nclass MPS(nn.Module):\n    \"\"\"\n    Tunable MPS model giving mapping from fixed-size data to output vector\n\n    Model works by first converting each 'pixel' (local data) to feature\n    vector via a simple embedding, then contracting embeddings with inputs\n    to each MPS cores. The resulting transition matrices are contracted\n    together along bond dimensions (i.e. hidden state spaces), with output\n    produced via an uncontracted edge of an additional output core.\n\n    MPS model permits many customizable behaviors, including custom\n    'routing' of MPS through the input, choice of boundary conditions\n    (meaning the model can act as a tensor train or a tensor ring),\n    GPU-friendly parallel evaluation, and an experimental mode to support\n    adaptive choice of bond dimensions based on singular value spectrum.\n\n    Args:\n        input_dim:       Number of 'pixels' in the input to the MPS\n        output_dim:      Size of the vectors output by MPS via output core\n        bond_dim:        Dimension of the 'bonds' connecting adjacent MPS\n                         cores, which act as hidden state spaces of the\n                         model. In adaptive mode, bond_dim instead\n                         specifies the maximum allowed bond dimension\n        feature_dim:     Size of the local feature spaces each pixel is\n                         embedded into (default: 2)\n        periodic_bc:     Whether MPS has periodic boundary conditions (i.e.\n                         is a tensor ring) or open boundary conditions\n                         (i.e. is a tensor train) (default: False)\n        parallel_eval:   Whether contraction of tensors is performed in a\n                         serial or parallel fashion. The former is less\n                         expensive for open boundary conditions, but\n                         parallelizes more poorly (default: False)\n        label_site:      Location in the MPS chain where output is placed\n                         (default: input_dim // 2)\n        path:            List specifying a path through the input data\n                         which MPS is 'routed' along. For example, choosing\n                         path=[0, 1, ..., input_dim-1] gives a standard\n                         in-order traversal (behavior when path=None), while\n                         path=[0, 2, ..., input_dim-1] specifies an MPS\n                         accepting input only from even-valued input pixels\n                         (default: None)\n        init_std:        Size of the Gaussian noise used in default\n                         near-identity initialization (default: 1e-9)\n        initializer:     Pytorch initializer for custom initialization of\n                         MPS cores, with None specifying default\n                         near-identity initialization (default: None)\n        use_bias:        Whether to use trainable bias matrices in MPS\n                         cores, which are initialized near the zero matrix\n                         (default: False)\n        adaptive_mode:   Whether MPS is trained with experimental adaptive\n                         bond dimensions selection (default: False)\n        cutoff:          Singular value cutoff controlling bond dimension\n                         adaptive selection (default: 1e-9)\n        merge_threshold: Number of inputs before adaptive MPS shifts its\n                         merge state once, with two shifts leading to the\n                         update of all bond dimensions (default: 2000)\n    \"\"\"\n\n    # TODO: Support arbitrary initializers\n    # TODO: Clean up the current treatment of initialization\n    # TODO: Resolve weirdness with fixed bias and initialization choice\n    # TODO: Add function to convert to canonical form\n    # TODO: Fix issue of no training when use_bias=False\n\n    def __init__(\n        self,\n        input_dim,\n        output_dim,\n        bond_dim,\n        feature_dim=2,\n        periodic_bc=False,\n        parallel_eval=False,\n        label_site=None,\n        path=None,\n        init_std=1e-9,\n        initializer=None,\n        use_bias=True,\n        adaptive_mode=False,\n        cutoff=1e-10,\n        merge_threshold=2000,\n    ):\n        super().__init__()\n\n        if label_site is None:\n            label_site = input_dim // 2\n        assert label_site >= 0 and label_site <= input_dim\n\n        # Using bias matrices in adaptive_mode is too complicated, so I'm\n        # disabling it here\n        if adaptive_mode:\n            use_bias = False\n\n        # Our MPS is made of two InputRegions separated by an OutputSite.\n        module_list = []\n        init_args = {\n            \"bond_str\": \"slri\",\n            \"shape\": [label_site, bond_dim, bond_dim, feature_dim],\n            \"init_method\": (\n                \"min_random_eye\" if adaptive_mode else \"random_zero\",\n                init_std,\n                output_dim,\n            ),\n        }\n\n        # The first input region\n        if label_site > 0:\n            tensor = init_tensor(**init_args)\n\n            module_list.append(InputRegion(tensor, use_bias=use_bias, fixed_bias=False))\n\n        # The output site\n        tensor = init_tensor(\n            shape=[output_dim, bond_dim, bond_dim],\n            bond_str=\"olr\",\n            init_method=(\n                \"min_random_eye\" if adaptive_mode else \"random_eye\",\n                init_std,\n                output_dim,\n            ),\n        )\n        module_list.append(OutputSite(tensor))\n\n        # The other input region\n        if label_site < input_dim:\n            init_args[\"shape\"] = [\n                input_dim - label_site,\n                bond_dim,\n                bond_dim,\n                feature_dim,\n            ]\n            tensor = init_tensor(**init_args)\n            module_list.append(InputRegion(tensor, use_bias=use_bias, fixed_bias=False))\n\n        # Initialize linear_region according to our adaptive_mode specification\n        if adaptive_mode:\n            self.linear_region = MergedLinearRegion(\n                module_list=module_list,\n                periodic_bc=periodic_bc,\n                parallel_eval=parallel_eval,\n                cutoff=cutoff,\n                merge_threshold=merge_threshold,\n            )\n\n            # Initialize the list of bond dimensions, which starts out constant\n            self.bond_list = bond_dim * torch.ones(input_dim + 2, dtype=torch.long)\n            if not periodic_bc:\n                self.bond_list[0], self.bond_list[-1] = 1, 1\n\n            # Initialize the list of singular values, which start out at -1\n            self.sv_list = -1.0 * torch.ones([input_dim + 2, bond_dim])\n\n        else:\n            self.linear_region = LinearRegion(\n                module_list=module_list,\n                periodic_bc=periodic_bc,\n                parallel_eval=parallel_eval,\n            )\n        assert len(self.linear_region) == input_dim\n\n        if path:\n            assert isinstance(path, (list, torch.Tensor))\n            assert len(path) == input_dim\n\n        # Set the rest of our MPS attributes\n        self.input_dim = input_dim\n        self.output_dim = output_dim\n        self.bond_dim = bond_dim\n        self.feature_dim = feature_dim\n        self.periodic_bc = periodic_bc\n        self.adaptive_mode = adaptive_mode\n        self.label_site = label_site\n        self.path = path\n        self.use_bias = use_bias\n        self.cutoff = cutoff\n        self.merge_threshold = merge_threshold\n        self.feature_map = None\n\n    def forward(self, input_data):\n        \"\"\"\n        Embed our data and pass it to an MPS with a single output site\n\n        Args:\n            input_data (Tensor): Input with shape [batch_size, input_dim] or\n                                 [batch_size, input_dim, feature_dim]. In the\n                                 former case, the data points are turned into\n                                 2D vectors using a default linear feature map.\n\n                                 When using a user-specified path, the size of\n                                 the second tensor mode need not exactly equal\n                                 input_dim, since the path variable is used to\n                                 slice a certain subregion of input_data. This\n                                 can be used to define multiple MPS 'strings',\n                                 which act on different parts of the input.\n        \"\"\"\n        # For custom paths, rearrange our input into the desired order\n        if self.path:\n            path_inputs = []\n            for site_num in self.path:\n                path_inputs.append(input_data[:, site_num])\n            input_data = torch.stack(path_inputs, dim=1)\n\n        # Embed our input data before feeding it into our linear region\n        input_data = self.embed_input(input_data)\n        output = self.linear_region(input_data)\n\n        # If we got a tuple as output, then use the last two entries to\n        # update our bond dimensions and singular values\n        if isinstance(output, tuple):\n            output, new_bonds, new_svs = output\n\n            assert len(new_bonds) == len(self.bond_list)\n            assert len(new_bonds) == len(new_svs)\n            for i, bond_dim in enumerate(new_bonds):\n                if bond_dim != -1:\n                    assert new_svs[i] is not -1\n                    self.bond_list[i] = bond_dim\n                    self.sv_list[i] = new_svs[i]\n\n        return output\n\n    def embed_input(self, input_data):\n        \"\"\"\n        Embed pixels of input_data into separate local feature spaces\n\n        Args:\n            input_data (Tensor):    Input with shape [batch_size, input_dim], or\n                                    [batch_size, input_dim, feature_dim]. In the\n                                    latter case, the data is assumed to already\n                                    be embedded, and is returned unchanged.\n\n        Returns:\n            embedded_data (Tensor): Input embedded into a tensor with shape\n                                    [batch_size, input_dim, feature_dim]\n        \"\"\"\n        assert len(input_data.shape) in [2, 3]\n        assert input_data.size(1) == self.input_dim\n\n        # If input already has a feature dimension, return it as is\n        if len(input_data.shape) == 3:\n            if input_data.size(2) != self.feature_dim:\n                raise ValueError(\n                    f\"input_data has wrong shape to be unembedded \"\n                    \"or pre-embedded data (input_data.shape = \"\n                    f\"{list(input_data.shape)}, feature_dim = {self.feature_dim})\"\n                )\n            return input_data\n\n        # Apply a custom embedding map if it has been defined by the user\n        if self.feature_map is not None:\n            f_map = self.feature_map\n            embedded_data = torch.stack(\n                [torch.stack([f_map(x) for x in batch]) for batch in input_data]\n            )\n\n            # Make sure our embedded input has the desired size\n            assert embedded_data.shape == torch.Size(\n                [input_data.size(0), self.input_dim, self.feature_dim]\n            )\n\n        # Otherwise, use a simple linear embedding map with feature_dim = 2\n        else:\n            if self.feature_dim != 2:\n                raise RuntimeError(\n                    f\"self.feature_dim = {self.feature_dim}, \"\n                    \"but default feature_map requires self.feature_dim = 2\"\n                )\n\n            embedded_data = torch.stack([input_data, 1 - input_data], dim=2)\n\n        return embedded_data\n\n    def register_feature_map(self, feature_map):\n        \"\"\"\n        Register a custom feature map to be used for embedding input data\n\n        Args:\n            feature_map (function): Takes a single scalar input datum and\n                                    returns an embedded representation of the\n                                    image. The output size of the function must\n                                    match self.feature_dim. If feature_map=None,\n                                    then the feature map will be reset to a\n                                    simple default linear embedding\n        \"\"\"\n        if feature_map is not None:\n            # Test to make sure feature_map outputs vector of proper size\n            out_shape = feature_map(torch.tensor(0)).shape\n            needed_shape = torch.Size([self.feature_dim])\n            if out_shape != needed_shape:\n                raise ValueError(\n                    \"Given feature_map returns values of size \"\n                    f\"{list(out_shape)}, but should return \"\n                    f\"values of size {list(needed_shape)}\"\n                )\n\n        self.feature_map = feature_map\n\n    def core_len(self):\n        \"\"\"\n        Returns the number of cores, which is at least the required input size\n        \"\"\"\n        return self.linear_region.core_len()\n\n    def __len__(self):\n        \"\"\"\n        Returns the number of input sites, which equals the input size\n        \"\"\"\n        return self.input_dim\n\n\nclass LinearRegion(nn.Module):\n    \"\"\"\n    List of modules which feeds input to each module and returns reduced output\n    \"\"\"\n\n    def __init__(\n        self, module_list, periodic_bc=False, parallel_eval=False, module_states=None\n    ):\n        # Check that module_list is a list whose entries are Pytorch modules\n        if not isinstance(module_list, list) or module_list is []:\n            raise ValueError(\"Input to LinearRegion must be nonempty list\")\n        for i, item in enumerate(module_list):\n            if not isinstance(item, nn.Module):\n                raise ValueError(\n                    \"Input items to LinearRegion must be PyTorch \"\n                    f\"Module instances, but item {i} is not\"\n                )\n        super().__init__()\n\n        # Wrap as a ModuleList for proper parameter registration\n        self.module_list = nn.ModuleList(module_list)\n        self.periodic_bc = periodic_bc\n        self.parallel_eval = parallel_eval\n\n    def forward(self, input_data):\n        \"\"\"\n        Contract input with list of MPS cores and return result as contractable\n\n        Args:\n            input_data (Tensor): Input with shape [batch_size, input_dim,\n                                                   feature_dim]\n        \"\"\"\n        # Check that input_data has the correct shape\n        assert len(input_data.shape) == 3\n        assert input_data.size(1) == len(self)\n        periodic_bc = self.periodic_bc\n        parallel_eval = self.parallel_eval\n        lin_bonds = [\"l\", \"r\"]\n\n        # Whether to move intermediate vectors to a GPU (fixes Issue #8)\n        to_cuda = input_data.is_cuda\n        device = f\"cuda:{input_data.get_device()}\" if to_cuda else \"cpu\"\n\n        # For each module, pull out the number of pixels needed and call that\n        # module's forward() method, putting the result in contractable_list\n        ind = 0\n        contractable_list = []\n        for module in self.module_list:\n            mod_len = len(module)\n            if mod_len == 1:\n                mod_input = input_data[:, ind]\n            else:\n                mod_input = input_data[:, ind : (ind + mod_len)]\n            ind += mod_len\n\n            contractable_list.append(module(mod_input))\n\n        # For periodic boundary conditions, reduce contractable_list and\n        # trace over the left and right indices to get our output\n        if periodic_bc:\n            contractable_list = ContractableList(contractable_list)\n            contractable = contractable_list.reduce(parallel_eval=True)\n\n            # Unpack the output (atomic) contractable\n            tensor, bond_str = contractable.tensor, contractable.bond_str\n            assert all(c in bond_str for c in lin_bonds)\n\n            # Build einsum string for the trace of tensor\n            in_str, out_str = \"\", \"\"\n            for c in bond_str:\n                if c in lin_bonds:\n                    in_str += \"l\"\n                else:\n                    in_str += c\n                    out_str += c\n            ein_str = in_str + \"->\" + out_str\n\n            # Return the trace over left and right indices\n            return torch.einsum(ein_str, [tensor])\n\n        # For open boundary conditions, add dummy edge vectors to\n        # contractable_list and reduce everything to get our output\n        else:\n            # Get the dimension of left and right bond indices\n            end_items = [contractable_list[i] for i in [0, -1]]\n            bond_strs = [item.bond_str for item in end_items]\n            bond_inds = [bs.index(c) for (bs, c) in zip(bond_strs, lin_bonds)]\n            bond_dims = [\n                item.tensor.size(ind) for (item, ind) in zip(end_items, bond_inds)\n            ]\n\n            # Build dummy end vectors and insert them at the ends of our list\n            end_vecs = [torch.zeros(dim).to(device) for dim in bond_dims]\n\n            for vec in end_vecs:\n                vec[0] = 1\n            contractable_list.insert(0, EdgeVec(end_vecs[0], is_left_vec=True))\n            contractable_list.append(EdgeVec(end_vecs[1], is_left_vec=False))\n\n            # Multiply together everything in contractable_list\n            contractable_list = ContractableList(contractable_list)\n            output = contractable_list.reduce(parallel_eval=parallel_eval)\n\n            return output.tensor\n\n    def core_len(self):\n        \"\"\"\n        Returns the number of cores, which is at least the required input size\n        \"\"\"\n        return sum([module.core_len() for module in self.module_list])\n\n    def __len__(self):\n        \"\"\"\n        Returns the number of input sites, which is the required input size\n        \"\"\"\n        return sum([len(module) for module in self.module_list])\n\n\nclass MergedLinearRegion(LinearRegion):\n    \"\"\"\n    Dynamic variant of LinearRegion that periodically rearranges its submodules\n    \"\"\"\n\n    def __init__(\n        self,\n        module_list,\n        periodic_bc=False,\n        parallel_eval=False,\n        cutoff=1e-10,\n        merge_threshold=2000,\n    ):\n        # Initialize a LinearRegion with our given module_list\n        super().__init__(module_list, periodic_bc, parallel_eval)\n\n        # Initialize attributes self.module_list_0 and self.module_list_1\n        # using the unmerged self.module_list, then redefine the latter in\n        # terms of one of the former lists\n        self.offset = 0\n        self._merge(offset=self.offset)\n        self._merge(offset=(self.offset + 1) % 2)\n        self.module_list = getattr(self, f\"module_list_{self.offset}\")\n\n        # Initialize variables used during switching\n        self.input_counter = 0\n        self.merge_threshold = merge_threshold\n        self.cutoff = cutoff\n\n    def forward(self, input_data):\n        \"\"\"\n        Contract input with list of MPS cores and return result as contractable\n\n        MergedLinearRegion keeps an input counter of the number of inputs, and\n        when this exceeds its merge threshold, triggers an unmerging and\n        remerging of its parameter tensors.\n\n        Args:\n            input_data (Tensor): Input with shape [batch_size, input_dim,\n                                                   feature_dim]\n        \"\"\"\n        # If we've hit our threshold, flip the merge state of our tensors\n        if self.input_counter >= self.merge_threshold:\n            bond_list, sv_list = self._unmerge(cutoff=self.cutoff)\n            self.offset = (self.offset + 1) % 2\n            self._merge(offset=self.offset)\n            self.input_counter -= self.merge_threshold\n\n            # Point self.module_list to the appropriate merged module\n            self.module_list = getattr(self, f\"module_list_{self.offset}\")\n        else:\n            bond_list, sv_list = None, None\n\n        # Increment our counter and call the LinearRegion's forward method\n        self.input_counter += input_data.size(0)\n        output = super().forward(input_data)\n\n        # If we flipped our merge state, then return the bond_list and output\n        if bond_list:\n            return output, bond_list, sv_list\n        else:\n            return output\n\n    @torch.no_grad()\n    def _merge(self, offset):\n        \"\"\"\n        Convert unmerged modules in self.module_list to merged counterparts\n\n        Calling _merge (or _unmerge) directly can cause undefined behavior,\n        but see MergedLinearRegion.forward for intended use\n\n        This proceeds by first merging all unmerged cores internally, then\n        merging lone cores when possible during a second sweep\n        \"\"\"\n        assert offset in [0, 1]\n\n        unmerged_list = self.module_list\n\n        # Merge each core internally and add the results to midway_list\n        site_num = offset\n        merged_list = []\n        for core in unmerged_list:\n            assert not isinstance(core, MergedInput)\n            assert not isinstance(core, MergedOutput)\n\n            # Apply internal merging routine if our core supports it\n            if hasattr(core, \"_merge\"):\n                merged_list.extend(core._merge(offset=site_num % 2))\n            else:\n                merged_list.append(core)\n\n            site_num += core.core_len()\n\n        # Merge pairs of cores when possible (currently only with\n        # InputSites), making sure to respect the offset for merging.\n        while True:\n            mod_num, site_num = 0, 0\n            combined_list = []\n\n            while mod_num < len(merged_list) - 1:\n                left_core, right_core = merged_list[mod_num : mod_num + 2]\n                new_core = self.combine(left_core, right_core, merging=True)\n\n                # If cores aren't combinable, move our sliding window by 1\n                if new_core is None or offset != site_num % 2:\n                    combined_list.append(left_core)\n                    mod_num += 1\n                    site_num += left_core.core_len()\n\n                # If we get something new, move to the next distinct pair\n                else:\n                    assert (\n                        new_core.core_len()\n                        == left_core.core_len() + right_core.core_len()\n                    )\n                    combined_list.append(new_core)\n                    mod_num += 2\n                    site_num += new_core.core_len()\n\n                # Add the last core if there's nothing to merge it with\n                if mod_num == len(merged_list) - 1:\n                    combined_list.append(merged_list[mod_num])\n                    mod_num += 1\n\n            # We're finished when unmerged_list remains unchanged\n            if len(combined_list) == len(merged_list):\n                break\n            else:\n                merged_list = combined_list\n\n        # Finally, update the appropriate merged module list\n        list_name = f\"module_list_{offset}\"\n        # If the merged module list hasn't been set yet, initialize it\n        if not hasattr(self, list_name):\n            setattr(self, list_name, nn.ModuleList(merged_list))\n\n        # Otherwise, do an in-place update so that all tensors remain\n        # properly registered with whatever optimizer we use\n        else:\n            module_list = getattr(self, list_name)\n            assert len(module_list) == len(merged_list)\n            for i in range(len(module_list)):\n                assert module_list[i].tensor.shape == merged_list[i].tensor.shape\n                module_list[i].tensor[:] = merged_list[i].tensor\n\n    @torch.no_grad()\n    def _unmerge(self, cutoff=1e-10):\n        \"\"\"\n        Convert merged modules to unmerged counterparts\n\n        Calling _unmerge (or _merge) directly can cause undefined behavior,\n        but see MergedLinearRegion.forward for intended use\n\n        This proceeds by first unmerging all merged cores internally, then\n        combining lone cores where possible\n        \"\"\"\n        list_name = f\"module_list_{self.offset}\"\n        merged_list = getattr(self, list_name)\n\n        # Unmerge each core internally and add results to unmerged_list\n        unmerged_list, bond_list, sv_list = [], [-1], [-1]\n        for core in merged_list:\n\n            # Apply internal unmerging routine if our core supports it\n            if hasattr(core, \"_unmerge\"):\n                new_cores, new_bonds, new_svs = core._unmerge(cutoff)\n                unmerged_list.extend(new_cores)\n                bond_list.extend(new_bonds[1:])\n                sv_list.extend(new_svs[1:])\n            else:\n                assert not isinstance(core, InputRegion)\n                unmerged_list.append(core)\n                bond_list.append(-1)\n                sv_list.append(-1)\n\n        # Combine all combinable pairs of cores. This occurs in several\n        # passes, and for now acts nontrivially only on InputSite instances\n        while True:\n            mod_num = 0\n            combined_list = []\n\n            while mod_num < len(unmerged_list) - 1:\n                left_core, right_core = unmerged_list[mod_num : mod_num + 2]\n                new_core = self.combine(left_core, right_core, merging=False)\n\n                # If cores aren't combinable, move our sliding window by 1\n                if new_core is None:\n                    combined_list.append(left_core)\n                    mod_num += 1\n\n                # If we get something new, move to the next distinct pair\n                else:\n                    combined_list.append(new_core)\n                    mod_num += 2\n\n                # Add the last core if there's nothing to combine it with\n                if mod_num == len(unmerged_list) - 1:\n                    combined_list.append(unmerged_list[mod_num])\n                    mod_num += 1\n\n            # We're finished when unmerged_list remains unchanged\n            if len(combined_list) == len(unmerged_list):\n                break\n            else:\n                unmerged_list = combined_list\n\n        # Find the average (log) norm of all of our cores\n        log_norms = []\n        for core in unmerged_list:\n            log_norms.append([torch.log(norm) for norm in core.get_norm()])\n        log_scale = sum([sum(ns) for ns in log_norms])\n        log_scale /= sum([len(ns) for ns in log_norms])\n\n        # Now rescale all cores so that their norms are roughly equal\n        scales = [[torch.exp(log_scale - n) for n in ns] for ns in log_norms]\n        for core, these_scales in zip(unmerged_list, scales):\n            core.rescale_norm(these_scales)\n\n        # Add our unmerged module list as a new attribute and return\n        # the updated bond dimensions\n        self.module_list = nn.ModuleList(unmerged_list)\n        return bond_list, sv_list\n\n    def combine(self, left_core, right_core, merging):\n        \"\"\"\n        Combine a pair of cores into a new core using context-dependent rules\n\n        Depending on the types of left_core and right_core, along with whether\n        we're currently merging (merging=True) or unmerging (merging=False),\n        either return a new core, or None if no rule exists for this context\n        \"\"\"\n\n        # Combine an OutputSite with a stray InputSite, return a MergedOutput\n        if merging and (\n            (isinstance(left_core, OutputSite) and isinstance(right_core, InputSite))\n            or (isinstance(left_core, InputSite) and isinstance(right_core, OutputSite))\n        ):\n\n            left_site = isinstance(left_core, InputSite)\n            if left_site:\n                new_tensor = torch.einsum(\n                    \"lui,our->olri\", [left_core.tensor, right_core.tensor]\n                )\n            else:\n                new_tensor = torch.einsum(\n                    \"olu,uri->olri\", [left_core.tensor, right_core.tensor]\n                )\n            return MergedOutput(new_tensor, left_output=(not left_site))\n\n        # Combine an InputRegion with a stray InputSite, return an InputRegion\n        elif not merging and (\n            (isinstance(left_core, InputRegion) and isinstance(right_core, InputSite))\n            or (\n                isinstance(left_core, InputSite) and isinstance(right_core, InputRegion)\n            )\n        ):\n\n            left_site = isinstance(left_core, InputSite)\n            if left_site:\n                left_tensor = left_core.tensor.unsqueeze(0)\n                right_tensor = right_core.tensor\n            else:\n                left_tensor = left_core.tensor\n                right_tensor = right_core.tensor.unsqueeze(0)\n\n            assert left_tensor.shape[1:] == right_tensor.shape[1:]\n            new_tensor = torch.cat([left_tensor, right_tensor])\n\n            return InputRegion(new_tensor)\n\n        # If this situation doesn't belong to the above cases, return None\n        else:\n            return None\n\n    def core_len(self):\n        \"\"\"\n        Returns the number of cores, which is at least the required input size\n        \"\"\"\n        return sum([module.core_len() for module in self.module_list])\n\n    def __len__(self):\n        \"\"\"\n        Returns the number of input sites, which is the required input size\n        \"\"\"\n        return sum([len(module) for module in self.module_list])\n\n\nclass InputRegion(nn.Module):\n    \"\"\"\n    Contiguous region of MPS input cores, associated with bond_str = 'slri'\n    \"\"\"\n\n    def __init__(\n        self, tensor, use_bias=True, fixed_bias=True, bias_mat=None, ephemeral=False\n    ):\n        super().__init__()\n\n        # Make sure tensor has correct size and the component mats are square\n        assert len(tensor.shape) == 4\n        assert tensor.size(1) == tensor.size(2)\n        bond_dim = tensor.size(1)\n\n        # If we are using bias matrices, set those up here\n        if use_bias:\n            assert bias_mat is None or isinstance(bias_mat, torch.Tensor)\n            bias_mat = (\n                torch.eye(bond_dim).unsqueeze(0) if bias_mat is None else bias_mat\n            )\n\n            bias_modes = len(list(bias_mat.shape))\n            assert bias_modes in [2, 3]\n            if bias_modes == 2:\n                bias_mat = bias_mat.unsqueeze(0)\n\n        # Register our tensors as a Pytorch Parameter or Tensor\n        if ephemeral:\n            self.register_buffer(name=\"tensor\", tensor=tensor.contiguous())\n            self.register_buffer(name=\"bias_mat\", tensor=bias_mat)\n        else:\n            self.register_parameter(\n                name=\"tensor\", param=nn.Parameter(tensor.contiguous())\n            )\n            if fixed_bias:\n                self.register_buffer(name=\"bias_mat\", tensor=bias_mat)\n            else:\n                self.register_parameter(name=\"bias_mat\", param=nn.Parameter(bias_mat))\n\n        self.use_bias = use_bias\n        self.fixed_bias = fixed_bias\n\n    def forward(self, input_data):\n        \"\"\"\n        Contract input with MPS cores and return result as a MatRegion\n\n        Args:\n            input_data (Tensor): Input with shape [batch_size, input_dim,\n                                                   feature_dim]\n        \"\"\"\n        # Check that input_data has the correct shape\n        tensor = self.tensor\n        assert len(input_data.shape) == 3\n        assert input_data.size(1) == len(self)\n        assert input_data.size(2) == tensor.size(3)\n\n        # Contract the input with our core tensor\n        mats = torch.einsum(\"slri,bsi->bslr\", [tensor, input_data])\n\n        # If we're using bias matrices, add those here\n        if self.use_bias:\n            bias_mat = self.bias_mat.unsqueeze(0)\n            mats = mats + bias_mat.expand_as(mats)\n\n        return MatRegion(mats)\n\n    def _merge(self, offset):\n        \"\"\"\n        Merge all pairs of neighboring cores and return a new list of cores\n\n        offset is either 0 or 1, which gives the first core at which we start\n        our merging. Depending on the length of our InputRegion, the output of\n        merge may have 1, 2, or 3 entries, with the majority of sites ending in\n        a MergedInput instance\n        \"\"\"\n        assert offset in [0, 1]\n        num_sites = self.core_len()\n        parity = num_sites % 2\n\n        # Cases with empty tensors might arise in recursion below\n        if num_sites == 0:\n            return [None]\n\n        # Simplify the problem into one where offset=0 and num_sites is even\n        if (offset, parity) == (1, 1):\n            out_list = [self[0], self[1:]._merge(offset=0)[0]]\n        elif (offset, parity) == (1, 0):\n            out_list = [self[0], self[1:-1]._merge(offset=0)[0], self[-1]]\n        elif (offset, parity) == (0, 1):\n            out_list = [self[:-1]._merge(offset=0)[0], self[-1]]\n\n        # The main case of interest, with no offset and an even number of sites\n        else:\n            tensor = self.tensor\n            even_cores, odd_cores = tensor[0::2], tensor[1::2]\n            assert len(even_cores) == len(odd_cores)\n\n            # Multiply all pairs of cores, keeping inputs separate\n            merged_cores = torch.einsum(\"slui,surj->slrij\", [even_cores, odd_cores])\n            out_list = [MergedInput(merged_cores)]\n\n        # Remove empty MergedInputs, which appear in very small InputRegions\n        return [x for x in out_list if x is not None]\n\n    def __getitem__(self, key):\n        \"\"\"\n        Returns an InputRegion instance sliced along the site index\n        \"\"\"\n        assert isinstance(key, int) or isinstance(key, slice)\n\n        if isinstance(key, slice):\n            return InputRegion(self.tensor[key])\n        else:\n            return InputSite(self.tensor[key])\n\n    def get_norm(self):\n        \"\"\"\n        Returns list of the norms of each core in InputRegion\n        \"\"\"\n        return [torch.norm(core) for core in self.tensor]\n\n    @torch.no_grad()\n    def rescale_norm(self, scale_list):\n        \"\"\"\n        Rescales the norm of each core by an amount specified in scale_list\n\n        For the i'th tensor defining a core in InputRegion, we rescale as\n        tensor_i <- scale_i * tensor_i, where scale_i = scale_list[i]\n        \"\"\"\n        assert len(scale_list) == len(self.tensor)\n\n        for core, scale in zip(self.tensor, scale_list):\n            core *= scale\n\n    def core_len(self):\n        return len(self)\n\n    def __len__(self):\n        return self.tensor.size(0)\n\n\nclass MergedInput(nn.Module):\n    \"\"\"\n    Contiguous region of merged MPS cores, each taking in a pair of input data\n\n    Since MergedInput arises after contracting together existing input cores,\n    a merged input tensor is required for initialization\n    \"\"\"\n\n    def __init__(self, tensor):\n        # Check that our input tensor has the correct shape\n        # bond_str = \"slrij\"\n        shape = tensor.shape\n        assert len(shape) == 5\n        assert shape[1] == shape[2]\n        assert shape[3] == shape[4]\n\n        super().__init__()\n\n        # Register our tensor as a Pytorch Parameter\n        self.register_parameter(name=\"tensor\", param=nn.Parameter(tensor.contiguous()))\n\n    def forward(self, input_data):\n        \"\"\"\n        Contract input with merged MPS cores and return result as a MatRegion\n\n        Args:\n            input_data (Tensor): Input with shape [batch_size, input_dim,\n                                 feature_dim], where input_dim must be even\n                                 (each merged core takes 2 inputs)\n        \"\"\"\n        # Check that input_data has the correct shape\n        tensor = self.tensor\n        assert len(input_data.shape) == 3\n        assert input_data.size(1) == len(self)\n        assert input_data.size(2) == tensor.size(3)\n        assert input_data.size(1) % 2 == 0\n\n        # Divide input_data into inputs living on even and on odd sites\n        inputs = [input_data[:, 0::2], input_data[:, 1::2]]\n\n        # Contract the odd (right-most) and even inputs with merged cores\n        tensor = torch.einsum(\"slrij,bsj->bslri\", [tensor, inputs[1]])\n        mats = torch.einsum(\"bslri,bsi->bslr\", [tensor, inputs[0]])\n\n        return MatRegion(mats)\n\n    def _unmerge(self, cutoff=1e-10):\n        \"\"\"\n        Separate the cores in our MergedInput and return an InputRegion\n\n        The length of the resultant InputRegion will be identical to our\n        original MergedInput (same number of inputs), but its core_len will\n        be doubled (twice as many individual cores)\n        \"\"\"\n        # bond_str = \"slrij\"\n        tensor = self.tensor\n        svd_string = \"lrij->lui,urj\"\n        max_D = tensor.size(1)\n\n        # Split every one of the cores into two and add them both to core_list\n        core_list, bond_list, sv_list = [], [-1], [-1]\n        for merged_core in tensor:\n            sv_vec = torch.empty(max_D)\n            left_core, right_core, bond_dim = svd_flex(\n                merged_core, svd_string, max_D, cutoff, sv_vec=sv_vec\n            )\n\n            core_list += [left_core, right_core]\n            bond_list += [bond_dim, -1]\n            sv_list += [sv_vec, -1]\n\n        # Collate the split cores into one tensor and return as an InputRegion\n        tensor = torch.stack(core_list)\n        return [InputRegion(tensor)], bond_list, sv_list\n\n    def get_norm(self):\n        \"\"\"\n        Returns list of the norm of each core in MergedInput\n        \"\"\"\n        return [torch.norm(core) for core in self.tensor]\n\n    @torch.no_grad()\n    def rescale_norm(self, scale_list):\n        \"\"\"\n        Rescales the norm of each core by an amount specified in scale_list\n\n        For the i'th tensor defining a core in MergedInput, we rescale as\n        tensor_i <- scale_i * tensor_i, where scale_i = scale_list[i]\n        \"\"\"\n        assert len(scale_list) == len(self.tensor)\n\n        for core, scale in zip(self.tensor, scale_list):\n            core *= scale\n\n    def core_len(self):\n        return len(self)\n\n    def __len__(self):\n        \"\"\"\n        Returns the number of input sites, which is twice the number of cores\n        \"\"\"\n        return 2 * self.tensor.size(0)\n\n\nclass InputSite(nn.Module):\n    \"\"\"\n    A single MPS core which takes in a single input datum, bond_str = 'lri'\n    \"\"\"\n\n    def __init__(self, tensor):\n        super().__init__()\n        # Register our tensor as a Pytorch Parameter\n        self.register_parameter(name=\"tensor\", param=nn.Parameter(tensor.contiguous()))\n\n    def forward(self, input_data):\n        \"\"\"\n        Contract input with MPS core and return result as a SingleMat\n\n        Args:\n            input_data (Tensor): Input with shape [batch_size, feature_dim]\n        \"\"\"\n        # Check that input_data has the correct shape\n        tensor = self.tensor\n        assert len(input_data.shape) == 2\n        assert input_data.size(1) == tensor.size(2)\n\n        # Contract the input with our core tensor\n        mat = torch.einsum(\"lri,bi->blr\", [tensor, input_data])\n\n        return SingleMat(mat)\n\n    def get_norm(self):\n        \"\"\"\n        Returns the norm of our core tensor, wrapped as a singleton list\n        \"\"\"\n        return [torch.norm(self.tensor)]\n\n    @torch.no_grad()\n    def rescale_norm(self, scale):\n        \"\"\"\n        Rescales the norm of our core by a factor of input `scale`\n        \"\"\"\n        if isinstance(scale, list):\n            assert len(scale) == 1\n            scale = scale[0]\n\n        self.tensor *= scale\n\n    def core_len(self):\n        return 1\n\n    def __len__(self):\n        return 1\n\n\nclass OutputSite(nn.Module):\n    \"\"\"\n    A single MPS core with no input and a single output index, bond_str = 'olr'\n    \"\"\"\n\n    def __init__(self, tensor):\n        super().__init__()\n        # Register our tensor as a Pytorch Parameter\n        self.register_parameter(name=\"tensor\", param=nn.Parameter(tensor.contiguous()))\n\n    def forward(self, input_data):\n        \"\"\"\n        Return the OutputSite wrapped as an OutputCore contractable\n        \"\"\"\n        return OutputCore(self.tensor)\n\n    def get_norm(self):\n        \"\"\"\n        Returns the norm of our core tensor, wrapped as a singleton list\n        \"\"\"\n        return [torch.norm(self.tensor)]\n\n    @torch.no_grad()\n    def rescale_norm(self, scale):\n        \"\"\"\n        Rescales the norm of our core by a factor of input `scale`\n        \"\"\"\n        if isinstance(scale, list):\n            assert len(scale) == 1\n            scale = scale[0]\n\n        self.tensor *= scale\n\n    def core_len(self):\n        return 1\n\n    def __len__(self):\n        return 0\n\n\nclass MergedOutput(nn.Module):\n    \"\"\"\n    Merged MPS core taking in one input datum and returning an output vector\n\n    Since MergedOutput arises after contracting together an existing input and\n    output core, an already-merged tensor is required for initialization\n\n    Args:\n        tensor (Tensor):    Value that our merged core is initialized to\n        left_output (bool): Specifies if the output core is on the left side of\n                            the input core (True), or on the right (False)\n    \"\"\"\n\n    def __init__(self, tensor, left_output):\n        # Check that our input tensor has the correct shape\n        # bond_str = \"olri\"\n        assert len(tensor.shape) == 4\n        super().__init__()\n\n        # Register our tensor as a Pytorch Parameter\n        self.register_parameter(name=\"tensor\", param=nn.Parameter(tensor.contiguous()))\n        self.left_output = left_output\n\n    def forward(self, input_data):\n        \"\"\"\n        Contract input with input index of core and return an OutputCore\n\n        Args:\n            input_data (Tensor): Input with shape [batch_size, feature_dim]\n        \"\"\"\n        # Check that input_data has the correct shape\n        tensor = self.tensor\n        assert len(input_data.shape) == 2\n        assert input_data.size(1) == tensor.size(3)\n\n        # Contract the input with our core tensor\n        tensor = torch.einsum(\"olri,bi->bolr\", [tensor, input_data])\n\n        return OutputCore(tensor)\n\n    def _unmerge(self, cutoff=1e-10):\n        \"\"\"\n        Split our MergedOutput into an OutputSite and an InputSite\n\n        The non-zero entries of our tensors are dynamically sized according to\n        the SVD cutoff, but will generally be padded with zeros to give the\n        new index a regular size.\n        \"\"\"\n        # bond_str = \"olri\"\n        tensor = self.tensor\n        left_output = self.left_output\n        if left_output:\n            svd_string = \"olri->olu,uri\"\n            max_D = tensor.size(2)\n            sv_vec = torch.empty(max_D)\n\n            output_core, input_core, bond_dim = svd_flex(\n                tensor, svd_string, max_D, cutoff, sv_vec=sv_vec\n            )\n            return (\n                [OutputSite(output_core), InputSite(input_core)],\n                [-1, bond_dim, -1],\n                [-1, sv_vec, -1],\n            )\n\n        else:\n            svd_string = \"olri->our,lui\"\n            max_D = tensor.size(1)\n            sv_vec = torch.empty(max_D)\n\n            output_core, input_core, bond_dim = svd_flex(\n                tensor, svd_string, max_D, cutoff, sv_vec=sv_vec\n            )\n            return (\n                [InputSite(input_core), OutputSite(output_core)],\n                [-1, bond_dim, -1],\n                [-1, sv_vec, -1],\n            )\n\n    def get_norm(self):\n        \"\"\"\n        Returns the norm of our core tensor, wrapped as a singleton list\n        \"\"\"\n        return [torch.norm(self.tensor)]\n\n    @torch.no_grad()\n    def rescale_norm(self, scale):\n        \"\"\"\n        Rescales the norm of our core by a factor of input `scale`\n        \"\"\"\n        if isinstance(scale, list):\n            assert len(scale) == 1\n            scale = scale[0]\n\n        self.tensor *= scale\n\n    def core_len(self):\n        return 2\n\n    def __len__(self):\n        return 1\n\n\nclass InitialVector(nn.Module):\n    \"\"\"\n    Vector of ones and zeros to act as initial vector within the MPS\n\n    By default the initial vector is chosen to be all ones, but if fill_dim is\n    specified then only the first fill_dim entries are set to one, with the\n    rest zero.\n\n    If fixed_vec is False, then the initial vector will be registered as a\n    trainable model parameter.\n    \"\"\"\n\n    def __init__(self, bond_dim, fill_dim=None, fixed_vec=True, is_left_vec=True):\n        super().__init__()\n\n        vec = torch.ones(bond_dim)\n        if fill_dim is not None:\n            assert fill_dim >= 0 and fill_dim <= bond_dim\n            vec[fill_dim:] = 0\n\n        if fixed_vec:\n            vec.requires_grad = False\n            self.register_buffer(name=\"vec\", tensor=vec)\n        else:\n            vec.requires_grad = True\n            self.register_parameter(name=\"vec\", param=nn.Parameter(vec))\n\n        assert isinstance(is_left_vec, bool)\n        self.is_left_vec = is_left_vec\n\n    def forward(self):\n        \"\"\"\n        Return our initial vector wrapped as an EdgeVec contractable\n        \"\"\"\n        return EdgeVec(self.vec, self.is_left_vec)\n\n    def core_len(self):\n        return 1\n\n    def __len__(self):\n        return 0\n\n\nclass TerminalOutput(nn.Module):\n    \"\"\"\n    Output matrix at end of chain to transmute virtual state into output vector\n\n    By default, a fixed rectangular identity matrix with shape\n    [bond_dim, output_dim] will be used as a state transducer. If fixed_mat is\n    False, then the matrix will be registered as a trainable model parameter.\n    \"\"\"\n\n    def __init__(self, bond_dim, output_dim, fixed_mat=False, is_left_mat=False):\n        super().__init__()\n\n        # I don't have a nice initialization scheme for a non-injective fixed\n        # state transducer, so just throw an error if that's needed\n        if fixed_mat and output_dim > bond_dim:\n            raise ValueError(\n                \"With fixed_mat=True, TerminalOutput currently \"\n                \"only supports initialization for bond_dim >= \"\n                \"output_dim, but here bond_dim=\"\n                f\"{bond_dim} and output_dim={output_dim}\"\n            )\n\n        # Initialize the matrix and register it appropriately\n        mat = torch.eye(bond_dim, output_dim)\n        if fixed_mat:\n            mat.requires_grad = False\n            self.register_buffer(name=\"mat\", tensor=mat)\n        else:\n            # Add some noise to help with training\n            mat = mat + torch.randn_like(mat) / bond_dim\n\n            mat.requires_grad = True\n            self.register_parameter(name=\"mat\", param=nn.Parameter(mat))\n\n        assert isinstance(is_left_mat, bool)\n        self.is_left_mat = is_left_mat\n\n    def forward(self):\n        \"\"\"\n        Return our terminal matrix wrapped as an OutputMat contractable\n        \"\"\"\n        return OutputMat(self.mat, self.is_left_mat)\n\n    def core_len(self):\n        return 1\n\n    def __len__(self):\n        return 0","metadata":{"execution":{"iopub.status.busy":"2021-06-27T13:49:08.512433Z","iopub.execute_input":"2021-06-27T13:49:08.512801Z","iopub.status.idle":"2021-06-27T13:49:08.687548Z","shell.execute_reply.started":"2021-06-27T13:49:08.512769Z","shell.execute_reply":"2021-06-27T13:49:08.686543Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"mps = MPS(\n    input_dim=512 ** 2,\n    output_dim=4,\n    bond_dim=bond_dim,\n    adaptive_mode=adaptive_mode,\n    periodic_bc=periodic_bc,\n)","metadata":{"execution":{"iopub.status.busy":"2021-06-27T13:49:31.147973Z","iopub.execute_input":"2021-06-27T13:49:31.148382Z","iopub.status.idle":"2021-06-27T13:49:31.367840Z","shell.execute_reply.started":"2021-06-27T13:49:31.148347Z","shell.execute_reply":"2021-06-27T13:49:31.366809Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"loss_fun = torch.nn.CrossEntropyLoss()\noptimizer = torch.optim.Adam(mps.parameters(), lr=learn_rate, weight_decay=l2_reg)","metadata":{"execution":{"iopub.status.busy":"2021-06-27T13:49:37.798229Z","iopub.execute_input":"2021-06-27T13:49:37.798871Z","iopub.status.idle":"2021-06-27T13:49:37.804659Z","shell.execute_reply.started":"2021-06-27T13:49:37.798822Z","shell.execute_reply":"2021-06-27T13:49:37.803870Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(f\"Maximum MPS bond dimension = {bond_dim}\")\nprint(f\" * {'Adaptive' if adaptive_mode else 'Fixed'} bond dimensions\")\nprint(f\" * {'Periodic' if periodic_bc else 'Open'} boundary conditions\")\nprint(f\"Using Adam w/ learning rate = {learn_rate:.1e}\")\nif l2_reg > 0:\n    print(f\" * L2 regularization = {l2_reg:.2e}\")\nprint()","metadata":{"execution":{"iopub.status.busy":"2021-06-27T13:49:42.021824Z","iopub.execute_input":"2021-06-27T13:49:42.022479Z","iopub.status.idle":"2021-06-27T13:49:42.031015Z","shell.execute_reply.started":"2021-06-27T13:49:42.022413Z","shell.execute_reply":"2021-06-27T13:49:42.029814Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')","metadata":{"execution":{"iopub.status.busy":"2021-06-27T13:49:46.339364Z","iopub.execute_input":"2021-06-27T13:49:46.339790Z","iopub.status.idle":"2021-06-27T13:49:46.345122Z","shell.execute_reply.started":"2021-06-27T13:49:46.339755Z","shell.execute_reply":"2021-06-27T13:49:46.344185Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%bash\nconda install -c conda-forge gdcm -y","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"num_train = 500\nnum_test = 200","metadata":{"execution":{"iopub.status.busy":"2021-06-27T13:49:49.701155Z","iopub.execute_input":"2021-06-27T13:49:49.701742Z","iopub.status.idle":"2021-06-27T13:49:49.706942Z","shell.execute_reply.started":"2021-06-27T13:49:49.701690Z","shell.execute_reply":"2021-06-27T13:49:49.705880Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Put COVID data into dataloaders\nsamplers = {\n    \"train\": torch.utils.data.SubsetRandomSampler(range(num_train)),\n    \"test\": torch.utils.data.SubsetRandomSampler(range(num_test)),\n}\nloaders = {\n    name: torch.utils.data.DataLoader(\n        dataset, batch_size=batch_size, sampler=samplers[name], drop_last=True\n    )\n    for (name, dataset) in [(\"train\", transformed_train_dataset), (\"test\", transformed_test_dataset)]\n}\nnum_batches = {\n    name: total_num // batch_size\n    for (name, total_num) in [(\"train\", num_train), (\"test\", num_test)]\n}","metadata":{"execution":{"iopub.status.busy":"2021-06-27T13:49:52.718034Z","iopub.execute_input":"2021-06-27T13:49:52.718440Z","iopub.status.idle":"2021-06-27T13:49:52.725881Z","shell.execute_reply.started":"2021-06-27T13:49:52.718408Z","shell.execute_reply":"2021-06-27T13:49:52.725082Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"transformed_train_dataset","metadata":{"execution":{"iopub.status.busy":"2021-06-27T17:10:20.275055Z","iopub.execute_input":"2021-06-27T17:10:20.275527Z","iopub.status.idle":"2021-06-27T17:10:20.281507Z","shell.execute_reply.started":"2021-06-27T17:10:20.275487Z","shell.execute_reply":"2021-06-27T17:10:20.280746Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Let's start training!\nfor epoch_num in range(1, num_epochs + 1):\n    running_loss = 0.0\n    running_acc = 0.0\n\n    for inputs, labels in transformed_train_dataset:\n        inputs, labels = inputs.view([batch_size, 512 ** 2]), labels.data\n        #inputs = inputs.data, labels.data\n\n        # Call our MPS to get logit scores and predictions\n        scores = mps(inputs)\n        _, preds = torch.max(scores, 1)\n\n        # Compute the loss and accuracy, add them to the running totals\n        loss = loss_fun(scores, labels)\n        with torch.no_grad():\n            accuracy = torch.sum(preds == labels).item() / batch_size\n            running_loss += loss\n            running_acc += accuracy\n\n        # Backpropagate and update parameters\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n\n    print(f\"### Epoch {epoch_num} ###\")\n    print(f\"Average loss:           {running_loss / num_batches['train']:.4f}\")\n    print(f\"Average train accuracy: {running_acc / num_batches['train']:.4f}\")\n\n    # Evaluate accuracy of MPS classifier on the validation set\n    with torch.no_grad():\n        running_acc = 0.0\n\n        for inputs, labels in loaders[\"test\"]:\n            inputs, labels = inputs.view([batch_size, 1512 ** 2]), labels.data\n\n            # Call our MPS to get logit scores and predictions\n            scores = mps(inputs)\n            _, preds = torch.max(scores, 1)\n            running_acc += torch.sum(preds == labels).item() / batch_size\n\n    print(f\"Validation accuracy:          {running_acc / num_batches['valid']:.4f}\")\n    print(f\"Runtime so far:         {int(time.time()-start_time)} sec\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-27T17:00:10.926394Z","iopub.execute_input":"2021-06-27T17:00:10.926783Z","iopub.status.idle":"2021-06-27T17:00:11.388142Z","shell.execute_reply.started":"2021-06-27T17:00:10.926748Z","shell.execute_reply":"2021-06-27T17:00:11.386373Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}