modules.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "from collections import Counter\n",
    "import numpy as np\n",
    "from numbers import Number\n",
    "from functools import reduce\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.autograd import Variable\n",
    "from copy import deepcopy\n",
    "\n",
    "import sys, os\n",
    "sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "from pytorch_net.util import get_activation, get_activation_noise, init_weight, init_bias, init_module_weights, init_module_bias, to_np_array, to_Variable, zero_grad_hook, ACTIVATION_LIST\n",
    "from pytorch_net.util import standardize_symbolic_expression, get_param_name_list, get_variable_name_list, get_list_DL, get_coeffs_tree, snap, unsnap\n",
    "AVAILABLE_REG = [\"L1\", \"L2\", \"param\"]\n",
    "Default_Activation = \"linear\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Register all layer types:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_Layer(layer_type, input_size, output_size, W_init = None, b_init = None, settings = {}, is_cuda = False):\n",
    "    \"\"\"Obtain layer from specifications.\"\"\"\n",
    "    if layer_type == \"Simple_Layer\":\n",
    "        layer = Simple_Layer(input_size=input_size,\n",
    "                             output_size=output_size,\n",
    "                             W_init=W_init,\n",
    "                             b_init=b_init,\n",
    "                             settings=settings,\n",
    "                             is_cuda=is_cuda,\n",
    "                            )\n",
    "    elif layer_type == \"SuperNet_Layer\":\n",
    "        layer = SuperNet_Layer(input_size=input_size,\n",
    "                               output_size=output_size,\n",
    "                               W_init=W_init,\n",
    "                               b_init=b_init,\n",
    "                               settings=settings,\n",
    "                               is_cuda=is_cuda,\n",
    "                              )\n",
    "    elif layer_type == \"Symbolic_Layer\":\n",
    "        layer = Symbolic_Layer(input_size=input_size,\n",
    "                               output_size=output_size,\n",
    "                               W_init=W_init,\n",
    "                               b_init=b_init,\n",
    "                               settings=settings,\n",
    "                               is_cuda=is_cuda,\n",
    "                              )\n",
    "    elif layer_type == \"Utility_Layer\":\n",
    "        layer = Utility_Layer(input_size=input_size,\n",
    "                              output_size=output_size,\n",
    "                              settings=settings,\n",
    "                              is_cuda=is_cuda,\n",
    "                             )\n",
    "    else:\n",
    "        raise Exception(\"layer_type '{}' not recognized!\".format(layer_type))\n",
    "    return layer\n",
    "\n",
    "\n",
    "def load_layer_dict(layer_dict, layer_type, is_cuda=False):\n",
    "    \"\"\"Load layer from layer_dict.\"\"\"\n",
    "    new_layer = get_Layer(layer_type=\"Symbolic_Layer\",\n",
    "                          input_size=layer_dict[\"input_size\"],\n",
    "                          output_size=layer_dict[\"output_size\"],\n",
    "                          W_init=layer_dict[\"weights\"],\n",
    "                          b_init=layer_dict[\"bias\"],\n",
    "                          settings=layer_dict[\"settings\"],\n",
    "                          is_cuda=is_cuda,\n",
    "                         )\n",
    "    return new_layer\n",
    "\n",
    "\n",
    "def Simple_2_Symbolic(simple_layer, settings={}, mode=\"normal\", prefix=\"\"):\n",
    "    \"\"\"Transform Simple Layer to Symbolic Layer.\"\"\"\n",
    "    from sympy import Symbol, Function\n",
    "    input_size = simple_layer.input_size\n",
    "    output_size = simple_layer.output_size\n",
    "    symbolic_expression = []\n",
    "    W_core, b_core = simple_layer.get_weights_bias()\n",
    "    W_init = {}\n",
    "\n",
    "    if mode == \"normal\":\n",
    "        for j in range(output_size):\n",
    "            expression = 0\n",
    "            if W_core is not None:\n",
    "                for i in range(input_size):\n",
    "                    expression += Symbol(\"{0}W{1}{2}\".format(prefix, i, j)) * Symbol(\"x{0}\".format(i))\n",
    "                    W_init[\"{0}W{1}{2}\".format(prefix, i, j)] = W_core[i, j]\n",
    "            if b_core is not None:\n",
    "                expression += Symbol(\"{0}b{1}\".format(prefix, j))\n",
    "                W_init[\"{0}b{1}\".format(prefix, j)] = b_core[j]\n",
    "            if \"activation\" in simple_layer.settings:\n",
    "                activation_name = simple_layer.settings[\"activation\"]\n",
    "            elif \"activation\" in settings:\n",
    "                activation_name = settings[\"activation\"]\n",
    "            else:\n",
    "                activation_name = Default_Activation\n",
    "            if activation_name != \"linear\":\n",
    "                expression = Function(activation_name)(expression)\n",
    "            symbolic_expression.append(expression)\n",
    "    elif mode == \"separable\":\n",
    "        vector_p, vector_q = snap(W_core, \"separable\")[0]\n",
    "        for j in range(output_size):\n",
    "            expression = 0\n",
    "            for i in range(input_size):\n",
    "                expression += Symbol(\"x{}\".format(i)) * Symbol(\"{0}p{1}\".format(prefix, i)) * Symbol(\"{0}q{1}\".format(prefix, j))\n",
    "                W_init[\"{0}p{1}\".format(prefix, i)] = vector_p[i]\n",
    "            expression += Symbol(\"{0}b{1}\".format(prefix, j))\n",
    "            W_init[\"{0}q{1}\".format(prefix, j)] = vector_q[j]\n",
    "            W_init[\"{0}b{1}\".format(prefix, j)] = b_core[j]\n",
    "            if \"activation\" in simple_layer.settings:\n",
    "                activation_name = simple_layer.settings[\"activation\"]\n",
    "            elif \"activation\" in settings:\n",
    "                activation_name = settings[\"activation\"]\n",
    "            else:\n",
    "                activation_name = Default_Activation        \n",
    "            if activation_name != \"linear\":\n",
    "                expression = Function(activation_name)(expression)\n",
    "            symbolic_expression.append(expression)   \n",
    "\n",
    "    return get_Layer(layer_type=\"Symbolic_Layer\",\n",
    "                     input_size=input_size,\n",
    "                     output_size=output_size,\n",
    "                     W_init=W_init,\n",
    "                     b_init=None,\n",
    "                     settings={\"symbolic_expression\": str(symbolic_expression)},\n",
    "                     is_cuda=simple_layer.is_cuda,\n",
    "                    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simple Layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "SyntaxError",
     "evalue": "invalid syntax (<ipython-input-1-557b432823cf>, line 67)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;36m  File \u001b[0;32m\"<ipython-input-1-557b432823cf>\"\u001b[0;36m, line \u001b[0;32m67\u001b[0m\n\u001b[0;31m    activation = self.settings[]\u001b[0m\n\u001b[0m                               ^\u001b[0m\n\u001b[0;31mSyntaxError\u001b[0m\u001b[0;31m:\u001b[0m invalid syntax\n"
     ]
    }
   ],
   "source": [
    "class Simple_Layer(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        input_size,\n",
    "        output_size,\n",
    "        W_init=None,     # initialization for weights\n",
    "        b_init=None,     # initialization for bias\n",
    "        settings={},     # Other settings that are relevant to this specific layer\n",
    "        is_cuda=False,\n",
    "        ):\n",
    "        # Firstly, must perform this step:\n",
    "        super(Simple_Layer, self).__init__()\n",
    "        # Saving the attribuites:\n",
    "        if isinstance(input_size, tuple):\n",
    "            self.input_size = reduce(lambda x, y: x * y, input_size)\n",
    "            self.input_size_original = input_size\n",
    "        else:\n",
    "            self.input_size = input_size\n",
    "        if isinstance(output_size, tuple):\n",
    "            self.output_size = reduce(lambda x, y: x * y, output_size)\n",
    "            self.output_size_original = output_size\n",
    "        else:\n",
    "            self.output_size = output_size\n",
    "\n",
    "        self.is_cuda = is_cuda\n",
    "        self.device = torch.device(self.is_cuda if isinstance(self.is_cuda, str) else \"cuda\" if self.is_cuda else \"cpu\")\n",
    "        self.settings = settings\n",
    "\n",
    "        # Other attributes that are specific to this layer:\n",
    "        self.activation = settings[\"activation\"] if \"activation\" in settings else Default_Activation\n",
    "        self.weight_on = settings[\"weight_on\"] if \"weight_on\" in settings else True\n",
    "        self.bias_on = settings[\"bias_on\"] if \"bias_on\" in settings else True\n",
    "        self.reg_on = settings[\"reg_on\"] if \"reg_on\" in settings else True\n",
    "\n",
    "        # Define the learnable parameters in the module (use any name you like). \n",
    "        # use nn.Parameter() so that the parameters is registered in the module and can be gradient-updated:\n",
    "        # self.W_init, self.b_init can be a numpy array, or a string like \"glorot-normal\":\n",
    "        if self.weight_on:\n",
    "            self.W_core = nn.Parameter(torch.randn(self.input_size, self.output_size))\n",
    "            init_weight(self.W_core, init=W_init)\n",
    "        if self.bias_on:\n",
    "            self.b_core = nn.Parameter(torch.zeros(self.output_size))\n",
    "            init_bias(self.b_core, init=b_init)\n",
    "        # Dropout:\n",
    "        if \"dropout_rate\" in settings:\n",
    "            self.dropout = nn.Dropout(p=settings[\"dropout_rate\"])\n",
    "        self.set_cuda(is_cuda)\n",
    "\n",
    "        # Initialize parameter freeze if stipulated:\n",
    "        if \"snap_dict\" in self.settings:\n",
    "            # Clear snapping if either self.weight_on is False or self.bias_on is False\n",
    "            pop_snapping = []\n",
    "            for pos, idx in self.settings[\"snap_dict\"]:\n",
    "                if (self.weight_on is False and pos == \"weight\") or (self.bias_on is False and pos == \"bias\"):\n",
    "                    pop_snapping.append((pos, idx))\n",
    "            for key in pop_snapping:\n",
    "                self.settings[\"snap_dict\"].pop(key)\n",
    "        \n",
    "            # Initialize freeze:\n",
    "            self.snap_dict = self.settings[\"snap_dict\"]\n",
    "            self.initialize_param_freeze(update_values=True)\n",
    "        else:\n",
    "            self.snap_dict = {}\n",
    "\n",
    "\n",
    "    def __repr__(self):\n",
    "        string = \"\"\n",
    "        if not self.weight_on:\n",
    "            string += \", weight_on=False\"\n",
    "        if not self.bias_on:\n",
    "            string += \", bias_on=False\"\n",
    "        if not self.reg_on:\n",
    "            string += \", reg_on=False\"\n",
    "        if \"dropout_rate\" in self.settings:\n",
    "            string += \", dropout_rate={}\".format(self.settings[\"dropout_rate\"])\n",
    "        if \"act_noise\" in self.settings:\n",
    "            string += \", act_noise={}\".format(self.settings[\"act_noise\"])\n",
    "        return 'Simple_Layer({}, \"{}\"{})'.format(self.output_size, self.activation, string)\n",
    "\n",
    "\n",
    "    def change(self, target, new_property):\n",
    "        if target == \"weight\":\n",
    "            if self.weight_on:\n",
    "                old_property = \"on\"\n",
    "                if new_property == \"off\":\n",
    "                    self.settings[\"weight_on\"] = False\n",
    "                    self.weight_on = False\n",
    "                    delattr(self, \"W_core\")\n",
    "            else:\n",
    "                old_property = \"off\"\n",
    "                if new_property == \"on\":\n",
    "                    self.settings.pop(\"weight_on\")\n",
    "                    self.weight_on = True\n",
    "                    self.W_core = nn.Parameter(torch.randn(self.input_size, self.output_size))\n",
    "                    init_weight(self.W_core, init=None)\n",
    "\n",
    "        elif target == \"bias\":\n",
    "            if self.bias_on:\n",
    "                old_property = \"on\"\n",
    "                if new_property == \"off\":\n",
    "                    self.settings[\"bias_on\"] = False\n",
    "                    self.bias_on = False\n",
    "                    delattr(self, \"b_core\")\n",
    "            else:\n",
    "                old_property = \"off\"\n",
    "                if new_property == \"on\":\n",
    "                    self.settings.pop(\"bias_on\")\n",
    "                    self.bias_on = True\n",
    "                    self.b_core = nn.Parameter(torch.zeros(self.output_size))\n",
    "                    init_bias(self.b_core, init=None)\n",
    "                \n",
    "        elif target == \"activation\":\n",
    "            old_property = self.settings[\"activation\"]\n",
    "            self.settings[\"activation\"] = new_property\n",
    "            self.activation = self.settings[\"activation\"]  \n",
    "        else:\n",
    "            raise Exception(\"target can only be activation!\")\n",
    "        return old_property\n",
    "\n",
    "\n",
    "    @property\n",
    "    def struct_param(self):\n",
    "        output_size = self.output_size_original if hasattr(self, \"output_size_original\") else self.output_size\n",
    "        if len(self.snap_dict) > 0:\n",
    "            self.settings[\"snap_dict\"] = self.snap_dict\n",
    "        return [output_size, \"Simple_Layer\", self.settings]\n",
    "\n",
    "\n",
    "    @property\n",
    "    def layer_dict(self):\n",
    "        input_size = self.input_size_original if hasattr(self, \"input_size_original\") else self.input_size\n",
    "        output_size = self.output_size_original if hasattr(self, \"output_size_original\") else self.output_size\n",
    "        Layer_dict =  {\n",
    "            \"input_size\": input_size,\n",
    "            \"output_size\": output_size,\n",
    "            \"settings\": self.settings,\n",
    "        }\n",
    "        if len(self.snap_dict) > 0:\n",
    "            Layer_dict[\"settings\"][\"snap_dict\"] = self.snap_dict\n",
    "        Layer_dict[\"weights\"], Layer_dict[\"bias\"] = self.get_weights_bias()\n",
    "        return Layer_dict\n",
    "\n",
    "\n",
    "    @property\n",
    "    def DL(self):\n",
    "        non_snapped_list = []\n",
    "        snapped_list = []\n",
    "        # Weights:\n",
    "        if self.weight_on:\n",
    "            shape = self.W_core.shape\n",
    "            for i in range(shape[0]):\n",
    "                for j in range(shape[1]):\n",
    "                    if (\"weight\", (i, j)) in self.snap_dict:\n",
    "                        snapped_list.append(self.snap_dict[(\"weight\", (i, j))][\"new_value\"])\n",
    "                    else:\n",
    "                        non_snapped_list.append(to_np_array(self.W_core[i, j]))\n",
    "        # Bias:\n",
    "        if self.bias_on:\n",
    "            for i in range(len(self.b_core)):\n",
    "                if (\"bias\", i) in self.snap_dict:\n",
    "                    snapped_list.append(self.snap_dict[(\"bias\", i)][\"new_value\"])\n",
    "                else:\n",
    "                    non_snapped_list.append(to_np_array(self.b_core[i]))\n",
    "        return get_list_DL(snapped_list, \"snapped\") + get_list_DL(non_snapped_list, \"non-snapped\")\n",
    "\n",
    "\n",
    "    def load_layer_dict(self, layer_dict):\n",
    "        new_layer = load_layer_dict(layer_dict, \"Simple_Layer\", self.is_cuda)\n",
    "        self.__dict__.update(new_layer.__dict__)\n",
    "\n",
    "\n",
    "    def forward(self, input, p_dict=None):\n",
    "        output = input\n",
    "        if hasattr(self, \"input_size_original\"):\n",
    "            output = output.view(-1, self.input_size)\n",
    "        # Dropout:\n",
    "        if hasattr(self, \"dropout\"):\n",
    "            output = self.dropout(output)\n",
    "\n",
    "        # Perform dot(X, W) + b:\n",
    "        if self.weight_on:\n",
    "            output = torch.matmul(output, self.W_core)\n",
    "        if self.bias_on:\n",
    "            output = output + self.b_core\n",
    "        \n",
    "        # If p_dict is not None, update the first neuron's activation according to p_dict:\n",
    "        if p_dict is not None:\n",
    "            p_dict = p_dict.view(-1)\n",
    "            if len(p_dict) == 2:\n",
    "                output_0 = output[:,:1] * p_dict[1] + p_dict[0]\n",
    "            elif len(p_dict) == 1:\n",
    "                output_0 = output[:,:1] + p_dict[0]\n",
    "            else:\n",
    "                raise\n",
    "            if output.size(1) > 1:\n",
    "                output = torch.cat([output_0, output[:,1:]], 1)\n",
    "            else:\n",
    "                output = output_0\n",
    "\n",
    "        # Perform activation function:\n",
    "        output = get_activation(self.activation)(output)\n",
    "\n",
    "        # Add activation noise:\n",
    "        if \"act_noise\" in self.settings:\n",
    "            output = get_activation_noise(self.settings[\"act_noise\"])(output)\n",
    "\n",
    "        if hasattr(self, \"output_size_original\"):\n",
    "            output = output.view(*((-1,) + self.output_size_original))\n",
    "        assert output.size(0) == input.size(0), \"output_size {0} must have same length as input_size {1}. Check shape!\".format(output.size(0), input.size(0))\n",
    "        return output\n",
    "\n",
    "\n",
    "    def prune_output_neurons(self, neuron_ids):\n",
    "        if not isinstance(neuron_ids, list):\n",
    "            neuron_ids = [neuron_ids]\n",
    "        preserved_ids = torch.LongTensor(np.array(list(set(range(self.output_size)) - set(neuron_ids)))).to(self.device)\n",
    "        if self.weight_on:\n",
    "            self.W_core = nn.Parameter(self.W_core.data[:, preserved_ids])\n",
    "            self.output_size = self.W_core.shape[1]\n",
    "        if self.bias_on:\n",
    "            self.b_core = nn.Parameter(self.b_core.data[preserved_ids])\n",
    "            self.output_size = self.b_core.shape[0]\n",
    "    \n",
    "    \n",
    "    def prune_input_neurons(self, neuron_ids):\n",
    "        if self.weight_on:\n",
    "            if not isinstance(neuron_ids, list):\n",
    "                neuron_ids = [neuron_ids]\n",
    "            preserved_ids = torch.LongTensor(np.array(list(set(range(self.input_size)) - set(neuron_ids))))\n",
    "            self.W_core = nn.Parameter(self.W_core.data[preserved_ids, :])\n",
    "            self.input_size = self.W_core.size(0)\n",
    "        else:\n",
    "            print(\"Cannot shrink input neurons since weight_on=False\")\n",
    "\n",
    "    \n",
    "    def add_output_neurons(self, num_neurons, mode=\"imitation\"):\n",
    "        if mode == \"imitation\":\n",
    "            if self.weight_on:\n",
    "                W_core_mean = to_np_array(self.W_core.mean())\n",
    "                W_core_std = to_np_array(self.W_core.std())\n",
    "                new_W_core = torch.randn(self.input_size, num_neurons) * W_core_std + W_core_mean\n",
    "            if self.bias_on:\n",
    "                b_core_mean = to_np_array(self.b_core.mean())\n",
    "                b_core_std = to_np_array(self.b_core.std())\n",
    "                new_b_core = torch.randn(num_neurons) * b_core_std + b_core_mean\n",
    "        elif mode == \"zeros\":\n",
    "            if self.weight_on:\n",
    "                new_W_core = torch.zeros(self.input_size, num_neurons)\n",
    "            if self.bias_on:\n",
    "                new_b_core = torch.zeros(num_neurons)\n",
    "        elif mode[0] == \"copy\":\n",
    "            neuron_id = mode[1]\n",
    "            if self.weight_on:\n",
    "                new_W_core = self.W_core[:, neuron_id: neuron_id + 1].detach().data\n",
    "            if self.bias_on:\n",
    "                new_b_core = self.b_core[neuron_id: neuron_id + 1].detach().data\n",
    "        else:\n",
    "            raise Exception(\"mode {0} not recognized!\".format(mode))\n",
    "        if self.weight_on:\n",
    "            self.W_core = nn.Parameter(torch.cat([self.W_core.data, new_W_core.to(self.device)], 1))\n",
    "        if self.bias_on:\n",
    "            self.b_core = nn.Parameter(torch.cat([self.b_core.data, new_b_core.to(self.device)], 0))\n",
    "        self.output_size += num_neurons\n",
    "        \n",
    "    \n",
    "    def add_input_neurons(self, num_neurons, mode=\"imitation\", position=\"end\"):\n",
    "        if self.weight_on:\n",
    "            if mode == \"imitation\":\n",
    "                W_core_mean = self.W_core.mean().item()\n",
    "                W_core_std = self.W_core.std().item()\n",
    "                new_W_core = torch.randn(num_neurons, self.output_size) * W_core_std + W_core_mean\n",
    "            elif mode == \"zeros\":\n",
    "                new_W_core = torch.zeros(num_neurons, self.output_size)\n",
    "            else:\n",
    "                raise Exception(\"mode {} not recognized!\".format(mode))\n",
    "            if position == \"end\":\n",
    "                self.W_core = nn.Parameter(torch.cat([self.W_core.data, new_W_core], 0))\n",
    "            else:\n",
    "                assert isinstance(position, Number)\n",
    "                self.W_core = nn.Parameter(torch.cat([self.W_core.data[:position], new_W_core, self.W_core.data[position:]], 0))\n",
    "            self.input_size += num_neurons\n",
    "        else:\n",
    "            print(\"Cannot add input neurons since weight_on=False\")\n",
    "        \n",
    "\n",
    "    def standardize(self, mode=\"b_mean_zero\"):\n",
    "        if mode == \"b_mean_zero\":\n",
    "            if self.bias_on:\n",
    "                b_mean = to_np_array(self.b_core.mean())\n",
    "                self.b_core.data.copy_(self.b_core.data - b_mean)\n",
    "        else:\n",
    "            raise Exception(\"mode {0} not recognized!\".format(mode))\n",
    "\n",
    "\n",
    "    def simplify(self, mode=\"snap\", excluded_idx=[], top=1, **kwargs):\n",
    "        def get_idx_list(key_list, input_size, output_size, weight_on):\n",
    "            \"\"\"Transform (pos, true_idx) list to idx list.\"\"\"\n",
    "            idx_list = []\n",
    "            for pos, true_idx in key_list:\n",
    "                if pos == \"weight\":\n",
    "                    assert self.weight_on is True\n",
    "                    idx_list.append(true_idx[0] * output_size + true_idx[1])\n",
    "                elif pos == \"bias\":\n",
    "                    assert self.bias_on is True\n",
    "                    if weight_on:\n",
    "                        idx_list.append(true_idx + input_size * output_size)\n",
    "                    else:\n",
    "                        idx_list.append(true_idx)\n",
    "                else:\n",
    "                    raise\n",
    "            return sorted(idx_list)\n",
    "\n",
    "        def get_true_idx(idx, input_size, output_size, weight_on):\n",
    "            \"\"\"Get (pos, true_idx) from idx\"\"\"\n",
    "            if weight_on:\n",
    "                if idx < input_size * output_size:\n",
    "                    pos = \"weight\"\n",
    "                    true_idx = (int(idx / output_size), idx % output_size)\n",
    "                else:\n",
    "                    pos = \"bias\"\n",
    "                    true_idx = idx - input_size * output_size\n",
    "            else:\n",
    "                pos = \"bias\"\n",
    "                true_idx = idx\n",
    "            return pos, true_idx\n",
    "\n",
    "        if mode == \"snap\":\n",
    "            snap_mode = kwargs[\"snap_mode\"] if \"snap_mode\" in kwargs else \"integer\"\n",
    "            if snap_mode == \"unsnap\":\n",
    "                self.remove_param_freeze()\n",
    "                return [\"unsnap\"]\n",
    "            elif snap_mode == \"vector\":\n",
    "                return []\n",
    "            else:\n",
    "                # Identify the parameters to freeze:\n",
    "                param = []\n",
    "                if self.weight_on:\n",
    "                    param.append(to_np_array(self.W_core.view(-1)))\n",
    "                if self.bias_on:\n",
    "                    param.append(to_np_array(self.b_core.view(-1), full_reduce=False))\n",
    "                param = np.concatenate(param)\n",
    "                if \"snap_targets\" in kwargs and kwargs[\"snap_targets\"] is not None:\n",
    "                    snap_targets = kwargs[\"snap_targets\"]\n",
    "                    is_target_given = True\n",
    "                else:\n",
    "                    excluded_idx_combined = get_idx_list(set([element[0] for element in excluded_idx] + list(self.snap_dict.keys())), \n",
    "                                                         self.input_size, self.output_size, self.weight_on)\n",
    "                    snap_targets = snap(param, snap_mode=snap_mode, excluded_idx=excluded_idx_combined, top=top)\n",
    "                    is_target_given = False\n",
    "\n",
    "                info_list = []\n",
    "                for idx, new_value in snap_targets:\n",
    "                    if new_value is not None:\n",
    "                        if is_target_given:\n",
    "                            pos, true_idx = idx\n",
    "                            new_value = float(new_value)\n",
    "                        else:\n",
    "                            pos, true_idx = get_true_idx(idx, self.input_size, self.output_size, self.weight_on)\n",
    "                            new_value = new_value.astype(float)\n",
    "                        info_list.append(((pos, true_idx), new_value))\n",
    "                        if pos == \"weight\":\n",
    "                            new_W_core = self.W_core.data\n",
    "                            new_W_core[true_idx] = new_value\n",
    "                            self.W_core = nn.Parameter(new_W_core)\n",
    "                        elif pos == \"bias\":\n",
    "                            new_b_core = self.b_core.data\n",
    "                            new_b_core[true_idx] = new_value\n",
    "                            self.b_core = nn.Parameter(new_b_core)\n",
    "                        self.snap_dict[(pos, true_idx)] = {\"new_value\": new_value}\n",
    "                        self.initialize_param_freeze(update_values=False)\n",
    "        else:\n",
    "            raise Exception(\"mode {0} not recognized!\".format(mode))\n",
    "        return info_list\n",
    "\n",
    "\n",
    "    def initialize_param_freeze(self, update_values=True):\n",
    "        if update_values:\n",
    "            if self.weight_on:\n",
    "                new_W_core = self.W_core.data\n",
    "                for (pos, true_idx), item in self.snap_dict.items():\n",
    "                    if pos == \"weight\":\n",
    "                        new_W_core[true_idx] = item[\"new_value\"]\n",
    "                self.W_core = nn.Parameter(new_W_core)\n",
    "            if self.bias_on:\n",
    "                new_b_core = self.b_core.data\n",
    "                for (pos, true_idx), item in self.snap_dict.items():\n",
    "                    if pos == \"bias\":\n",
    "                        new_b_core[true_idx] = item[\"new_value\"]\n",
    "                self.b_core = nn.Parameter(new_b_core)\n",
    "        \n",
    "        # Initialize hook:\n",
    "        for pos, true_idx in self.snap_dict.keys():\n",
    "            hook_function = zero_grad_hook(true_idx)\n",
    "            if self.weight_on and pos == \"weight\":\n",
    "                h = self.W_core.register_hook(hook_function)\n",
    "            elif self.bias_on and pos == \"bias\":\n",
    "                h = self.b_core.register_hook(hook_function)\n",
    "    \n",
    "\n",
    "    def remove_param_freeze(self, index_list=None):\n",
    "        if index_list is None:\n",
    "            self.snap_dict = {}\n",
    "            self.settings.pop(\"snap_dict\")\n",
    "        else:\n",
    "            for key in index_list:\n",
    "                self.snap_dict.pop(key)\n",
    "            self.initialize_param_freeze(update_values=True)\n",
    "\n",
    "\n",
    "    def get_param_names(self, source):\n",
    "        if source == \"modules\":\n",
    "            if self.weight_on:\n",
    "                param_names = [\"W_core\"]\n",
    "            if self.bias_on:\n",
    "                param_names.append(\"b_core\")\n",
    "        if source == \"attention\":\n",
    "            param_names = []\n",
    "        return param_names\n",
    "\n",
    "\n",
    "    def get_weights_bias(self, is_grad=False):\n",
    "        if not is_grad:\n",
    "            W_core = deepcopy(to_np_array(self.W_core, full_reduce=False)) if self.weight_on else None\n",
    "            b_core = deepcopy(to_np_array(self.b_core, full_reduce=False)) if self.bias_on else None\n",
    "            return W_core, b_core\n",
    "        else:\n",
    "            W_grad = self.W_core.grad if self.weight_on else None\n",
    "            b_grad = self.b_core.grad if self.bias_on else None\n",
    "            W_grad = deepcopy(to_np_array(W_grad, full_reduce=False)) if W_grad is not None else None\n",
    "            b_grad = deepcopy(to_np_array(b_grad, full_reduce=False)) if b_grad is not None else None\n",
    "            return W_grad, b_grad\n",
    "\n",
    "    \n",
    "    def get_regularization(self, mode, source=[\"weight\", \"bias\"]):\n",
    "        if not isinstance(source, list):\n",
    "            source = [source]\n",
    "        reg = Variable(torch.FloatTensor(np.array([0])), requires_grad=False).to(self.device)\n",
    "        if self.reg_on:\n",
    "            for source_ele in source:\n",
    "                if self.weight_on:\n",
    "                    if source_ele == \"weight\":\n",
    "                        if mode == \"L1\":\n",
    "                            reg = reg + self.W_core.abs().sum()\n",
    "                        elif mode == \"L2\":\n",
    "                            reg = reg + (self.W_core ** 2).sum()\n",
    "                        elif mode in AVAILABLE_REG:\n",
    "                            pass\n",
    "                        else:\n",
    "                            raise Exception(\"mode '{}' not recognized!\".format(mode))\n",
    "                elif source_ele == \"bias\":\n",
    "                    if self.bias_on:\n",
    "                        if mode == \"L1\":\n",
    "                            reg = reg + self.b_core.abs().sum()\n",
    "                        elif mode == \"L2\":\n",
    "                            reg = reg + (self.b_core ** 2).sum()\n",
    "                        elif mode in AVAILABLE_REG:\n",
    "                            pass\n",
    "                        else:\n",
    "                            raise Exception(\"mode '{}' not recognized!\".format(mode))\n",
    "        return reg\n",
    "\n",
    "\n",
    "    def set_cuda(self, is_cuda):\n",
    "        if isinstance(is_cuda, str):\n",
    "            self.cuda(is_cuda)\n",
    "        else:\n",
    "            if is_cuda:\n",
    "                self.cuda()\n",
    "            else:\n",
    "                self.cpu()\n",
    "        self.is_cuda = is_cuda\n",
    "\n",
    "\n",
    "    def set_trainable(self, is_trainable):\n",
    "        if is_trainable:\n",
    "            if self.weight_on:\n",
    "                self.W_core.requires_grad = True\n",
    "            if self.bias_on:\n",
    "                self.b_core.requires_grad = True\n",
    "        else:\n",
    "            if self.weight_on:\n",
    "                self.W_core.requires_grad = False\n",
    "            if self.bias_on:\n",
    "                self.b_core.requires_grad = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Utility layers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Utility_Layer(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        input_size,\n",
    "        output_size,\n",
    "        settings={},\n",
    "        is_cuda=False,\n",
    "    ):\n",
    "        super(Utility_Layer, self).__init__()\n",
    "        self.input_size = input_size\n",
    "        self.output_size = output_size\n",
    "        self.is_cuda = is_cuda\n",
    "        self.settings = settings\n",
    "\n",
    "\n",
    "    def forward(self, input, **kwargs):\n",
    "        layer_type = self.settings[\"type\"]\n",
    "        if layer_type == \"reshape\":\n",
    "            return input.reshape(self.output_size)\n",
    "        elif layer_type == \"flatten\":\n",
    "            return input.view(-1)\n",
    "        else:\n",
    "            raise Exception(\"layer_type {} is not valid\".format(layer_type))\n",
    "\n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        model_dict = {}\n",
    "        model_dict[\"input_size\"] = self.input_size\n",
    "        model_dict[\"output_size\"] = self.output_size\n",
    "        model_dict[\"settings\"] = self.settings\n",
    "        return model_dict\n",
    "\n",
    "\n",
    "    @property\n",
    "    def struct_param(self):\n",
    "        return [self.output_size, \"Utility_Layer\", self.settings]\n",
    "\n",
    "\n",
    "    def get_regularization(self, mode, source=[\"weight\"], **kwargs):\n",
    "        reg = to_Variable([0], is_cuda=self.is_cuda)\n",
    "        return reg\n",
    "\n",
    "\n",
    "    def set_trainable(self, is_trainable):\n",
    "        pass\n",
    "\n",
    "\n",
    "    def set_cuda(self, is_cuda):\n",
    "        self.is_cuda = is_cuda"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Symbolic Layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Symbolic_Layer(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        input_size,\n",
    "        output_size,\n",
    "        W_init=None,\n",
    "        b_init=None,\n",
    "        settings={},\n",
    "        is_cuda=False,\n",
    "        ):\n",
    "        super(Symbolic_Layer, self).__init__()\n",
    "        from sympy.parsing.sympy_parser import parse_expr\n",
    "        self.input_size = input_size\n",
    "        self.output_size = output_size\n",
    "        self.W_init = W_init # Here we use W_init to represent all parameter initial values\n",
    "        self.is_cuda = is_cuda\n",
    "        self.is_numerical = False\n",
    "        self.set_symbolic_expression(str(settings[\"symbolic_expression\"]), p_init=self.W_init)\n",
    "\n",
    "\n",
    "    @property\n",
    "    def layer_dict(self):\n",
    "        return {\n",
    "            \"input_size\": self.input_size,\n",
    "            \"output_size\": self.output_size,\n",
    "            \"weights\": self.get_param_dict(),\n",
    "            \"bias\": None,\n",
    "            \"settings\": {\"symbolic_expression\": str(self.symbolic_expression)},\n",
    "        }\n",
    "\n",
    "\n",
    "    @property\n",
    "    def struct_param(self):\n",
    "        return [self.output_size, \"Symbolic_Layer\", {\"symbolic_expression\": str(self.symbolic_expression)}]\n",
    "\n",
    "\n",
    "    @property\n",
    "    def settings(self):\n",
    "        return {\"symbolic_expression\": str(self.symbolic_expression)}\n",
    "    \n",
    "    \n",
    "    @property\n",
    "    def activation(self):\n",
    "        activation_list = []\n",
    "        for expression in self.symbolic_expression:\n",
    "            if hasattr(expression.func, \"name\") and expression.func.name in ACTIVATION_LIST:\n",
    "                act_ele = expression.func.name\n",
    "            else:\n",
    "                act_ele = \"linear\"\n",
    "            activation_list.append(act_ele)\n",
    "        if len(Counter(activation_list)) > 1:\n",
    "            return \"linear\"\n",
    "        else:\n",
    "            return list(Counter(activation_list).keys())[0]\n",
    "    \n",
    "    \n",
    "    def change(self, target, new_property):\n",
    "        from sympy import Add\n",
    "        from sympy.utilities.lambdify import implemented_function\n",
    "        assert target == \"activation\"\n",
    "        prev_activation = self.activation\n",
    "        activation = new_property\n",
    "        if activation != prev_activation:\n",
    "            if prev_activation == \"linear\":\n",
    "                f = implemented_function(activation, get_activation(activation))\n",
    "                new_symbolic_expression = [f(expression) for expression in self.symbolic_expression]\n",
    "                self.set_symbolic_expression(new_symbolic_expression)\n",
    "            else:\n",
    "                if activation != \"linear\":\n",
    "                    f = implemented_function(activation, get_activation(activation))\n",
    "                    new_symbolic_expression = [f(*expression.args) for expression in self.symbolic_expression]\n",
    "                else:\n",
    "                    new_symbolic_expression = [Add(*expression.args) for expression in self.symbolic_expression]\n",
    "                self.set_symbolic_expression(new_symbolic_expression)\n",
    "\n",
    "\n",
    "    @property\n",
    "    def numerical_expression(self):\n",
    "        from sympy import Symbol\n",
    "        \"\"\"Replace the parameter in symbolic_expression by their numerical values\"\"\"\n",
    "        substitution = [(Symbol(param_name), to_np_array(getattr(self, param_name))) for param_name in self.param_name_list]\n",
    "        return [expression.subs(substitution) for expression in self.symbolic_expression]\n",
    "    \n",
    "    \n",
    "    def set_numerical(self, is_numerical):\n",
    "        self.is_numerical = is_numerical\n",
    "\n",
    "\n",
    "    @property\n",
    "    def DL(self):\n",
    "        param_dict = self.get_param_dict()\n",
    "        expr_length, snapped_list = get_coeffs_tree(self.symbolic_expression, param_dict)\n",
    "        non_snapped_list = list(param_dict.values())\n",
    "        return get_list_DL(snapped_list, \"snapped\") + get_list_DL(non_snapped_list, \"non-snapped\")\n",
    "\n",
    "\n",
    "    def load_layer_dict(self, layer_dict):\n",
    "        new_layer = load_layer_dict(layer_dict, \"Symbolic_Layer\", self.is_cuda)\n",
    "        self.__dict__.update(new_layer.__dict__)\n",
    "\n",
    "\n",
    "    def prune_output_neurons(self, neuron_ids):\n",
    "        if not isinstance(neuron_ids, list):\n",
    "            neuron_ids = [neuron_ids]\n",
    "        variable_names = self.get_variable_name_list()\n",
    "        assert \"x\" not in variable_names, \"In order to prune output_neurons, 'x' cannot be in the symbolic_expression!\"\n",
    "        symbolic_expression = [expression for i, expression in enumerate(self.symbolic_expression) if i not in neuron_ids]\n",
    "        self.output_size = sum(self.get_expression_length(symbolic_expression))\n",
    "        self.set_symbolic_expression(symbolic_expression)\n",
    "    \n",
    "    \n",
    "    def standardize(self, mode=\"b_mean_zero\"):\n",
    "        from sympy import Function\n",
    "        if mode == \"b_mean_zero\":\n",
    "            param_dict = self.get_param_dict()\n",
    "            bias_list = []\n",
    "            for expression in self.symbolic_expression:\n",
    "                fun_name_list = self.get_function_name_list(expression)\n",
    "                if len(fun_name_list) == 1:\n",
    "                    expr = expression.args[0]\n",
    "                elif len(fun_name_list) == 0:\n",
    "                    expr = expression\n",
    "                else:\n",
    "                    raise Exception(\"There must be at most one activation function\")\n",
    "\n",
    "                vars_subs = {element: 0 for element in self.get_variable_name_list(expr)}\n",
    "                bias = expr.subs(vars_subs).subs(param_dict)\n",
    "                bias_list.append(bias)\n",
    "            bias_mean = np.mean(bias_list)\n",
    "\n",
    "            new_symbolic_expression = []\n",
    "            for expression in self.symbolic_expression:\n",
    "                fun_name_list = self.get_function_name_list(expression)\n",
    "                if len(fun_name_list) == 1:\n",
    "                    fun = Function(fun_name_list[0])\n",
    "                    expr = expression.args[0]\n",
    "                elif len(fun_name_list) == 0:\n",
    "                    expr = expression\n",
    "                else:\n",
    "                    raise Exception(\"There must be at most one activation function\")\n",
    "\n",
    "                expr = expr - bias_mean\n",
    "                if len(fun_name_list) == 1:\n",
    "                    expr = fun(expr)\n",
    "                new_symbolic_expression.append(expr)\n",
    "            self.set_symbolic_expression(new_symbolic_expression)            \n",
    "        else:\n",
    "            raise Exception(\"mode {0} not recognized!\".format(mode))\n",
    "    \n",
    "    \n",
    "    def init_with_p_dict(self, p_dict):\n",
    "        self.set_param_values(p_dict)\n",
    "\n",
    "\n",
    "    def init_bias_with_input(self, input, mode=\"std_sqrt\"):\n",
    "        pass\n",
    "\n",
    "\n",
    "    def get_param_name_list(self, symbolic_expression=None):\n",
    "        \"\"\"Get parameter names from a given symbolic expression\"\"\"\n",
    "        # Here in the Sympy_Net we assume that the input is always represented by Symbol(\"x\"), so \"x\" is excluded from param_name_list:\n",
    "        symbolic_expression = self.symbolic_expression if symbolic_expression is None else symbolic_expression\n",
    "        symbolic_expression = standardize_symbolic_expression(symbolic_expression)\n",
    "        return get_param_name_list(symbolic_expression)\n",
    "\n",
    "\n",
    "    def get_variable_name_list(self, symbolic_expression=None):\n",
    "        symbolic_expression = self.symbolic_expression if symbolic_expression is None else symbolic_expression\n",
    "        symbolic_expression = standardize_symbolic_expression(symbolic_expression)\n",
    "        return get_variable_name_list(symbolic_expression)\n",
    "\n",
    "\n",
    "    def get_function_name_list(self, symbolic_expression=None):\n",
    "        from sympy import Function\n",
    "        from sympy.utilities.lambdify import implemented_function\n",
    "        symbolic_expression = self.symbolic_expression if symbolic_expression is None else symbolic_expression\n",
    "        symbolic_expression = standardize_symbolic_expression(symbolic_expression)\n",
    "        function_name_list = list({element.func.__name__ for expression in symbolic_expression for element in expression.atoms(Function) if element.func.__name__ not in [\"linear\"]})\n",
    "        self.implemented_function = {}\n",
    "        for function_name in function_name_list:\n",
    "            try:\n",
    "                self.implemented_function[function_name] = implemented_function(Function(function_name), get_activation(function_name))\n",
    "            except:\n",
    "                pass\n",
    "        return function_name_list\n",
    "\n",
    "\n",
    "    def get_param_dict(self):\n",
    "        param_names = self.get_param_name_list()\n",
    "        return {param_name: to_np_array(getattr(self, param_name)) for param_name in param_names}\n",
    "\n",
    "\n",
    "    def set_param_values(self, new_param_values):\n",
    "        param_names = self.get_param_name_list()\n",
    "        for key, value in new_param_values.items():\n",
    "            if key in param_names:\n",
    "                if isinstance(value, Variable):\n",
    "                    value_core = value.data\n",
    "                elif isinstance(value, float) or isinstance(value, int):\n",
    "                    value_core = torch.FloatTensor(np.array([value]))\n",
    "                getattr(self, key).data.copy_(value_core.view(-1))\n",
    "\n",
    "\n",
    "    def get_weights_bias(self, is_grad=False):\n",
    "        if not is_grad:\n",
    "            return deepcopy(self.get_param_dict()), None\n",
    "        else:\n",
    "            param_names = self.get_param_name_list()\n",
    "            param_grad_dict = {}\n",
    "            for param_name in param_names:\n",
    "                grad = getattr(self, param_name).grad\n",
    "                if grad is not None:\n",
    "                    param_grad_dict[param_name] = grad.item()\n",
    "                else:\n",
    "                    param_grad_dict[param_name] = None\n",
    "            return param_grad_dict, None\n",
    "\n",
    "\n",
    "    def get_expression_length(self, symbolic_expression=None):\n",
    "        symbolic_expression = self.symbolic_expression if symbolic_expression is None else symbolic_expression\n",
    "        symbolic_expression = standardize_symbolic_expression(symbolic_expression)\n",
    "        length_list = []\n",
    "        for expression in symbolic_expression:\n",
    "            variable_list = self.get_variable_name_list([expression])\n",
    "            if \"x\" in variable_list:\n",
    "                assert len(variable_list) == 1, \"x cannot coexist with x1, x2, etc. in a single expression, since the dimension is not compatible\"\n",
    "                length = self.input_size\n",
    "            else:\n",
    "                length = 1\n",
    "            length_list.append(length)\n",
    "        return length_list\n",
    "\n",
    "\n",
    "    def set_symbolic_expression(self, symbolic_expression, p_init=None):\n",
    "        \"\"\"Set a new symbolic expression and update the parameterss\"\"\"\n",
    "        symbolic_expression = standardize_symbolic_expression(symbolic_expression)\n",
    "        assert sum(self.get_expression_length(symbolic_expression)) == self.output_size, \"symbolic_expression's combined output length must be equal to self.output_size!\"\n",
    "        self.old_param_name_list = self.get_param_name_list(self.symbolic_expression) if hasattr(self, \"symbolic_expression\") else []\n",
    "        self.symbolic_expression = symbolic_expression\n",
    "        self.param_name_list = self.get_param_name_list(symbolic_expression)\n",
    "        self.variable_name_list = self.get_variable_name_list(symbolic_expression)\n",
    "        self.get_function_name_list()\n",
    "        \n",
    "        # If the new expression has parameter names that did not appear in previous expression, create it:\n",
    "        for param_name in self.param_name_list:            \n",
    "            if not hasattr(self, param_name):\n",
    "                if p_init is not None:\n",
    "                    param_init = p_init[param_name] if param_name in p_init else None\n",
    "                else:\n",
    "                    param_init = None\n",
    "\n",
    "                if param_init is None:\n",
    "                    setattr(self, param_name, nn.Parameter(torch.randn(1)))\n",
    "                else:\n",
    "                    setattr(self, param_name, nn.Parameter(torch.FloatTensor(np.array([param_init]))))\n",
    "\n",
    "        # Delete class parameters that do not appear in the new symbolic expression:\n",
    "        param_name_to_delete = set(self.old_param_name_list) - set(self.param_name_list)\n",
    "        for param_name in param_name_to_delete:\n",
    "            delattr(self, param_name)\n",
    "\n",
    "        self.set_cuda(self.is_cuda)\n",
    "\n",
    "\n",
    "    def set_trainable(self, is_trainable):\n",
    "        param_name_list = self.get_param_name_list()\n",
    "        for param_name in param_name_list:\n",
    "            if is_trainable:\n",
    "                getattr(self, param_name).requires_grad = True\n",
    "            else:\n",
    "                getattr(self, param_name).requires_grad = False\n",
    "    \n",
    "    \n",
    "    def set_cuda(self, is_cuda):\n",
    "        if isinstance(is_cuda, str):\n",
    "            self.cuda(is_cuda)\n",
    "        else:\n",
    "            if is_cuda:\n",
    "                self.cuda()\n",
    "            else:\n",
    "                self.cpu()\n",
    "        self.is_cuda = is_cuda\n",
    "\n",
    "\n",
    "    def forward(self, input, p_dict=None):\n",
    "        from sympy import Symbol, lambdify, N\n",
    "        symbols = [Symbol(variable_name) for variable_name in self.variable_name_list]\n",
    "        if p_dict is None:\n",
    "            symbols = tuple(symbols + [Symbol(param_name) for param_name in self.param_name_list])  # Get symbolic variables\n",
    "        else:\n",
    "            symbols = tuple(symbols + [Symbol(param_name) for param_name in sorted(list(p_dict.keys())) if \"x\" not in param_name])\n",
    "        f_list = [lambdify(symbols, N(expression), torch) for expression in self.symbolic_expression]    # Obtain the lambda function f(x0, x1,..., param0, param1, ...)\n",
    "        # Obtain the data that will be fed into (x0, x1,..., param0, param1, ...):\n",
    "        variables_feed = []\n",
    "        for variable_name in self.variable_name_list:\n",
    "            if variable_name == \"x\":\n",
    "                variable_feed = input\n",
    "            else:\n",
    "                idx = int(variable_name[1:])\n",
    "                variable_feed = input[:, idx: idx + 1]\n",
    "            variables_feed.append(variable_feed)\n",
    "        if p_dict is None:\n",
    "            symbols_feed = variables_feed + [getattr(self, param_name) for param_name in self.param_name_list]\n",
    "        else:\n",
    "            symbols_feed = variables_feed + [p_dict[param_name] for param_name in sorted(list(p_dict.keys())) if \"x\" not in param_name]\n",
    "        output_list = []\n",
    "        for f in f_list:\n",
    "            output_ele = f(*symbols_feed)\n",
    "            if not isinstance(output_ele, Variable):\n",
    "                output_ele = to_Variable(torch.ones(input.shape[0], 1), is_cuda=self.is_cuda) * output_ele\n",
    "            elif len(output_ele.shape) < 2 or output_ele.shape[0] != input.shape[0]:\n",
    "                multiplier = to_Variable(torch.ones(input.shape[0], 1), is_cuda=self.is_cuda)\n",
    "                output_ele = output_ele * multiplier\n",
    "            output_list.append(output_ele)\n",
    "        return torch.cat(output_list, 1)\n",
    "\n",
    "\n",
    "    def simplify(self, mode=\"form\", **kwargs):\n",
    "        from sympy import simplify, Symbol\n",
    "        verbose = kwargs[\"verbose\"] if \"verbose\" in kwargs else 0\n",
    "        info_list = []\n",
    "        if not isinstance(mode, list):\n",
    "            mode = [mode]\n",
    "        for mode_ele in mode:\n",
    "            if mode_ele == \"form\":\n",
    "                prev_expression = self.symbolic_expression\n",
    "                new_expression = [simplify(expression) for expression in self.symbolic_expression]\n",
    "                self.set_symbolic_expression(new_expression)\n",
    "                if verbose > 0:\n",
    "                    print(\"Original expression:\\tsymbolic: {0}; \\t numerical: {1}\".format(prev_expression, self.numerical_expression))\n",
    "                    print(\"New  expression: \\tsymbolic: {0}; \\t numerical: {1}\".format(self.symbolic_expression, self.numerical_expression))                \n",
    "            elif mode_ele == \"snap\":\n",
    "                snap_mode = kwargs[\"snap_mode\"] if \"snap_mode\" in kwargs else \"integer\"\n",
    "                top = kwargs[\"top\"] if \"top\" in kwargs else 1\n",
    "                if snap_mode == \"unsnap\":\n",
    "                    unsnapped_expression, new_param_dict = unsnap(self.symbolic_expression, self.get_param_dict())\n",
    "                    self.set_symbolic_expression(unsnapped_expression, new_param_dict)\n",
    "                    info_list = info_list + [(mode_ele, snap_mode, unsnapped_expression)]\n",
    "                elif snap_mode in [\"vector\"]:\n",
    "                    param_dict = self.get_param_dict()\n",
    "                    param_dict_subs, new_param_dict = snap(param_dict, snap_mode=snap_mode, top=top)\n",
    "                    param_dict.update(new_param_dict)\n",
    "                    if verbose > 0:\n",
    "                        print(\"Original expression:\\tsymbolic: {0}; \\t numerical: {1}\".format(self.symbolic_expression, self.numerical_expression))\n",
    "                        print(\"Substitution:  \\t{0}, with value: {1}\".format(pprint_dict(param_dict_subs), pprint_dict(new_param_dict)))\n",
    "                    new_expression = [expression.subs(param_dict_subs) for expression in self.symbolic_expression]\n",
    "                    self.set_symbolic_expression(new_expression, param_dict)\n",
    "                    info_list = info_list + [(deepcopy(param_dict_subs), deepcopy(new_param_dict))]\n",
    "                    if verbose > 0:\n",
    "                        print(\"New  expression: \\tsymbolic: {0}; \\t numerical: {1}\".format(self.symbolic_expression, self.numerical_expression))\n",
    "                else:\n",
    "                    param_names = list(self.get_param_dict().keys())\n",
    "                    param_array = np.array(list(self.get_param_dict().values()))\n",
    "                    snap_targets = snap(param_array, snap_mode=snap_mode, top=top)\n",
    "                    if not (len(snap_targets) == 1 and snap_targets[0][1] is None):\n",
    "                        subs_targets = [(Symbol(param_names[idx]), new_value) for idx, new_value in snap_targets]\n",
    "                        prev_expression = self.symbolic_expression\n",
    "                        if verbose > 0:\n",
    "                            print(\"Original expression:\\tsymbolic: {0}; \\t numerical: {1}\".format(prev_expression, self.numerical_expression))\n",
    "                            print(\"Substitution:  \\t{0}\".format(subs_targets))\n",
    "                        new_expression = [expression.subs(subs_targets) for expression in self.symbolic_expression]\n",
    "                        self.set_symbolic_expression(new_expression)\n",
    "                        info_list = info_list + [(param_names[idx], new_value) for idx, new_value in snap_targets]\n",
    "                        if verbose > 0:\n",
    "                            print(\"New  expression: \\tsymbolic: {0}; \\t numerical: {1}\".format(self.symbolic_expression, self.numerical_expression))\n",
    "            elif mode_ele == \"pair_snap\":\n",
    "                if len(self.get_param_dict()) < 2:\n",
    "                    raise Exception(\"Less than 2 parameters. Cannot pair_snap!\")\n",
    "                else:\n",
    "                    def get_param_inverse_dict(Dict):\n",
    "                        inverse_dict = {}\n",
    "                        i = 0\n",
    "                        param_list = []\n",
    "                        for key, value in Dict.items():\n",
    "                            param_list.append(value)\n",
    "                            inverse_dict[i] = key\n",
    "                            i += 1\n",
    "                        return param_list, inverse_dict\n",
    "                    snap_mode = kwargs[\"snap_mode\"] if \"snap_mode\" in kwargs else \"integer\"\n",
    "                    top = kwargs[\"top\"] if \"top\" in kwargs else 1\n",
    "                    if snap_mode == \"integer\":\n",
    "                        snap_mode_whole = \"pair_integer\"\n",
    "                    elif snap_mode == \"rational\":\n",
    "                        snap_mode_whole = \"pair_rational\"\n",
    "                    else:\n",
    "                        raise Exception(\"snap_mode {} not recognized!\".format(snap_mode))\n",
    "                    param_list, inverse_dict = get_param_inverse_dict(self.get_param_dict())\n",
    "\n",
    "                    snap_targets = snap(param_list, snap_mode=snap_mode_whole, top=top)\n",
    "                    subs_targets = [(Symbol(inverse_dict[replace_id]), Symbol(inverse_dict[ref_id]) * ratio) for (replace_id, ref_id), ratio in snap_targets]\n",
    "                    prev_expression = self.symbolic_expression\n",
    "                    new_expression = [expression.subs(subs_targets) for expression in self.symbolic_expression]\n",
    "                    self.set_symbolic_expression(new_expression)\n",
    "                    info_list = info_list + [(inverse_dict[replace_id], \"{} * \".format(ratio) + inverse_dict[ref_id]) for (replace_id, ref_id), ratio in snap_targets]\n",
    "                    if verbose > 0:\n",
    "                        print(\"Original expression:\\tsymbolic: {}; \\t numerical: {}\".format(prev_expression, self.numerical_expression))\n",
    "                        print(\"Substitution:  \\t{}\".format(subs_targets))\n",
    "                        print(\"New  expression: \\tsymbolic: {}; \\t numerical: {}\".format(self.symbolic_expression, self.numerical_expression))\n",
    "            else:\n",
    "                raise Exception(\"mode {} not recognized!\".format(mode_ele))\n",
    "        return info_list\n",
    "\n",
    "\n",
    "    def get_regularization(self, mode, source=[\"weight\"], **kwargs):\n",
    "        reg = to_Variable([0], is_cuda=self.is_cuda)\n",
    "        if not isinstance(source, list):\n",
    "            source = [source]\n",
    "        param_list = [param for param in self.parameters()]\n",
    "        if len(param_list) > 0 and \"weight\" in source:\n",
    "            params = torch.cat(param_list)\n",
    "            scale_factor = kwargs[\"reg_scale_factor\"] if \"reg_scale_factor\" in kwargs else None\n",
    "            if mode == \"L1\":\n",
    "                if scale_factor is not None:\n",
    "                    reg_indi = (params * to_Variable(scale_factor, is_cuda = self.is_cuda)).abs().sum()\n",
    "                else:\n",
    "                    reg_indi = params.abs().sum()\n",
    "                reg = reg + reg_indi                        \n",
    "            elif mode == \"L2\":\n",
    "                if scale_factor is not None:\n",
    "                    reg_indi = torch.sum((params * to_Variable(scale_factor, is_cuda = self.is_cuda)) ** 2)\n",
    "                else:\n",
    "                    reg_indi = torch.sum(params ** 2)\n",
    "                reg = reg + reg_indi\n",
    "        return reg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SuperNet Layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SuperNet_Layer(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        input_size,\n",
    "        output_size,\n",
    "        W_init=None,     # initialization for weights\n",
    "        b_init=None,     # initialization for bias\n",
    "        settings={},\n",
    "        is_cuda=False,\n",
    "        ):\n",
    "        super(SuperNet_Layer, self).__init__()\n",
    "        # Saving the attribuites:\n",
    "        if isinstance(input_size, tuple):\n",
    "            self.input_size = reduce(lambda x, y: x * y, input_size)\n",
    "            self.input_size_original = input_size\n",
    "        else:\n",
    "            self.input_size = input_size\n",
    "        if isinstance(output_size, tuple):\n",
    "            self.output_size = reduce(lambda x, y: x * y, output_size)\n",
    "            self.output_size_original = output_size\n",
    "        else:\n",
    "            self.output_size = output_size\n",
    "        self.W_init = W_init\n",
    "        self.b_init = b_init\n",
    "        self.is_cuda = is_cuda\n",
    "        self.device = torch.device(self.is_cuda if isinstance(self.is_cuda, str) else \"cuda\" if self.is_cuda else \"cpu\")\n",
    "        \n",
    "        # Obtain additional initialization settings if provided:\n",
    "        self.W_available = settings[\"W_available\"] if \"W_available\" in settings else [\"dense\", \"Toeplitz\"]\n",
    "        self.b_available = settings[\"b_available\"] if \"b_available\" in settings else [\"dense\", \"None\"]\n",
    "        self.A_available = settings[\"A_available\"] if \"A_available\" in settings else [\"linear\", \"relu\"]\n",
    "        self.W_sig_init  = settings[\"W_sig_init\"] if \"W_sig_init\" in settings else None # initialization for the significance for the weights\n",
    "        self.b_sig_init  = settings[\"b_sig_init\"] if \"b_sig_init\" in settings else None # initialization for the significance for the bias\n",
    "        self.A_sig_init  = settings[\"A_sig_init\"] if \"A_sig_init\" in settings else None # initialization for the significance for the activations\n",
    "        for W_candidate in self.W_available:\n",
    "            if \"2D-in\" in W_candidate:\n",
    "                self.input_size_2D = settings[\"input_size_2D\"]\n",
    "            if \"2D-out\" in W_candidate:\n",
    "                self.output_size_2D = settings[\"output_size_2D\"]\n",
    "        for b_candidate in self.b_available:\n",
    "            if \"2D\" in b_candidate:\n",
    "                self.output_size_2D = settings[\"output_size_2D\"]\n",
    "        \n",
    "        # Initialize layer:\n",
    "        self.init_layer()\n",
    "        self.set_cuda(is_cuda)\n",
    "    \n",
    "    \n",
    "    @property\n",
    "    def settings(self):\n",
    "        layer_settings = {}\n",
    "        layer_settings[\"W_available\"] = deepcopy(self.W_available)\n",
    "        layer_settings[\"b_available\"] = deepcopy(self.b_available)\n",
    "        layer_settings[\"A_available\"] = deepcopy(self.A_available)\n",
    "        layer_settings[\"W_sig_init\"] = to_np_array(self.W_sig)\n",
    "        layer_settings[\"b_sig_init\"] = to_np_array(self.b_sig)\n",
    "        layer_settings[\"A_sig_init\"] = to_np_array(self.A_sig)\n",
    "        return layer_settings\n",
    "    \n",
    "    @property\n",
    "    def struct_param(self):\n",
    "        return [self.output_size, \"SuperNet_Layer\", self.settings]\n",
    "\n",
    "        \n",
    "    def init_layer(self):\n",
    "        self.W_layer_seed = nn.Parameter(torch.FloatTensor(np.random.randn(self.input_size, self.output_size)))\n",
    "        self.b_layer_seed = nn.Parameter(torch.zeros(self.output_size))\n",
    "        init_weight(self.W_layer_seed, init = self.W_init)\n",
    "        init_bias(self.b_layer_seed, init = self.b_init)\n",
    "        if \"arithmetic-series-in\" in self.W_available:\n",
    "            self.W_interval_j = nn.Parameter(torch.randn(self.output_size) / np.sqrt(self.input_size + self.output_size))\n",
    "        if \"arithmetic-series-out\" in self.W_available:\n",
    "            self.W_interval_i = nn.Parameter(torch.randn(self.input_size) / np.sqrt(self.input_size + self.output_size))\n",
    "        if \"arithmetic-series-2D-in\" in self.W_available:\n",
    "            self.W_mean_2D_in = nn.Parameter(torch.randn(self.output_size) / np.sqrt(self.input_size_2D[0] + self.input_size_2D[1] + self.output_size))\n",
    "            self.W_interval_2D_in = nn.Parameter(torch.randn(2, self.output_size) / np.sqrt(self.input_size_2D[0] + self.input_size_2D[1] + self.output_size))\n",
    "        if \"arithmetic-series-2D-out\" in self.W_available:\n",
    "            self.W_mean_2D_out = nn.Parameter(torch.randn(self.input_size) / np.sqrt(self.input_size + self.output_size_2D[0] + self.output_size_2D[1]))\n",
    "            self.W_interval_2D_out = nn.Parameter(torch.randn(2, self.input_size) / np.sqrt(self.input_size + self.output_size_2D[0] + self.output_size_2D[1]))\n",
    "        if \"arithmetic-series\" in self.b_available:\n",
    "            self.b_interval = nn.Parameter(torch.randn(1) / np.sqrt(self.output_size))\n",
    "        if \"arithmetic-series-2D\" in self.b_available:\n",
    "            self.b_mean_2D = nn.Parameter(torch.randn(1) / np.sqrt(self.output_size))\n",
    "            self.b_interval_2D = nn.Parameter(torch.randn(2) / np.sqrt(self.output_size_2D[0] + self.output_size_2D[1]))\n",
    "        \n",
    "        if self.W_sig_init is None:\n",
    "            self.W_sig = nn.Parameter(torch.zeros(len(self.W_available)))\n",
    "        else:\n",
    "            self.W_sig = nn.Parameter(torch.FloatTensor(self.W_sig_init))\n",
    "        if self.b_sig_init is None:\n",
    "            self.b_sig = nn.Parameter(torch.zeros(len(self.b_available)))\n",
    "        else:\n",
    "            self.b_sig = nn.Parameter(torch.FloatTensor(self.b_sig_init))\n",
    "        if self.A_sig_init is None:\n",
    "            self.A_sig = nn.Parameter(torch.zeros(len(self.A_available)))\n",
    "        else:\n",
    "            self.A_sig = nn.Parameter(torch.FloatTensor(self.A_sig_init))\n",
    "\n",
    "\n",
    "    def get_layers(self, source=[\"weight\", \"bias\"]):\n",
    "        \"\"\"All the different SuperNet layers are based on the same W_seed matrices. \n",
    "        For example, W_seed is based on the full self.W_layer_seed; \"Toeplitz\" is based on\n",
    "        the first row and first column of self.W_layer_seed to construct the Toeplitz matrix, etc.\n",
    "        \"\"\"\n",
    "        # Superimpose different weights:\n",
    "        if \"weight\" in source:\n",
    "            self.W_list = []\n",
    "            for weight_type in self.W_available:\n",
    "                if weight_type == \"dense\":\n",
    "                    W_layer = self.W_layer_seed\n",
    "                elif weight_type == \"Toeplitz\":\n",
    "                    W_layer_stacked = []\n",
    "                    if self.output_size > 1:\n",
    "                        inv_idx = torch.arange(self.output_size - 1, 0, -1).long().to(self.device)\n",
    "                        W_seed = torch.cat([self.W_layer_seed[0][inv_idx], self.W_layer_seed[:,0]])\n",
    "                    else:\n",
    "                        W_seed = self.W_layer_seed[:,0]\n",
    "                    for j in range(self.output_size):\n",
    "                        W_layer_stacked.append(W_seed[self.output_size - j - 1: self.output_size - j - 1 + self.input_size])\n",
    "                    W_layer = torch.stack(W_layer_stacked, 1)\n",
    "                elif weight_type == \"arithmetic-series-in\":\n",
    "                    mean_j = self.W_layer_seed.mean(0)\n",
    "                    idx_i = torch.FloatTensor(np.repeat(np.arange(self.input_size), self.output_size)).to(self.device)\n",
    "                    idx_j = torch.LongTensor(range(self.output_size) * self.input_size).to(self.device)\n",
    "                    offset = self.input_size / float(2) - 0.5\n",
    "                    W_layer = (mean_j[idx_j] + self.W_interval_j[idx_j] * Variable(idx_i - offset, requires_grad = False)).view(self.input_size, self.output_size)\n",
    "                elif weight_type == \"arithmetic-series-out\":\n",
    "                    mean_i = self.W_layer_seed.mean(1)\n",
    "                    idx_i = torch.LongTensor(np.repeat(np.arange(self.input_size), self.output_size)).to(self.device)\n",
    "                    idx_j = torch.FloatTensor(range(self.output_size) * self.input_size).to(self.device)\n",
    "                    offset = self.output_size / float(2) - 0.5\n",
    "                    W_layer = (mean_i[idx_i] + self.W_interval_i[idx_i] * Variable(idx_j - offset, requires_grad = False)).view(self.input_size, self.output_size)\n",
    "                elif weight_type == \"arithmetic-series-2D-in\":\n",
    "                    idx_i, idx_j, idx_k = np.meshgrid(range(self.input_size_2D[0]), range(self.input_size_2D[1]), range(self.output_size), indexing = \"ij\")\n",
    "                    idx_i = torch.from_numpy(idx_i).float().view(-1).to(self.device)\n",
    "                    idx_j = torch.from_numpy(idx_j).float().view(-1).to(self.device)\n",
    "                    idx_k = torch.from_numpy(idx_k).long().view(-1).to(self.device)\n",
    "                    offset_i = self.input_size_2D[0] / float(2) - 0.5\n",
    "                    offset_j = self.input_size_2D[1] / float(2) - 0.5\n",
    "                    W_layer = (self.W_mean_2D_in[idx_k] + \\\n",
    "                               self.W_interval_2D_in[:, idx_k][0] * Variable(idx_i - offset_i, requires_grad = False) + \\\n",
    "                               self.W_interval_2D_in[:, idx_k][1] * Variable(idx_j - offset_j, requires_grad = False)).view(self.input_size, self.output_size)\n",
    "                elif weight_type == \"arithmetic-series-2D-out\":\n",
    "                    idx_k, idx_i, idx_j = np.meshgrid(range(self.input_size), range(self.output_size_2D[0]), range(self.output_size_2D[1]), indexing = \"ij\")\n",
    "                    idx_k = torch.from_numpy(idx_k).long().view(-1).to(self.device)\n",
    "                    idx_i = torch.from_numpy(idx_i).float().view(-1).to(self.device)\n",
    "                    idx_j = torch.from_numpy(idx_j).float().view(-1).to(self.device)\n",
    "                    offset_i = self.output_size_2D[0] / float(2) - 0.5\n",
    "                    offset_j = self.output_size_2D[1] / float(2) - 0.5\n",
    "                    W_layer = (self.W_mean_2D_out[idx_k] + \\\n",
    "                               self.W_interval_2D_out[:, idx_k][0] * Variable(idx_i - offset_i, requires_grad = False) + \\\n",
    "                               self.W_interval_2D_out[:, idx_k][1] * Variable(idx_j - offset_j, requires_grad = False)).view(self.input_size, self.output_size)\n",
    "                else:\n",
    "                    raise Exception(\"weight_type '{0}' not recognized!\".format(weight_type))\n",
    "                self.W_list.append(W_layer)\n",
    "\n",
    "            if len(self.W_available) == 1:\n",
    "                self.W_core = W_layer\n",
    "            else:\n",
    "                self.W_list = torch.stack(self.W_list, dim = 2)\n",
    "                W_sig_softmax = nn.Softmax(dim = -1)(self.W_sig.unsqueeze(0))\n",
    "                self.W_core = torch.matmul(self.W_list, W_sig_softmax.transpose(1,0)).squeeze(2)\n",
    "    \n",
    "        # Superimpose different biases:\n",
    "        if \"bias\" in source:\n",
    "            self.b_list = []\n",
    "            for bias_type in self.b_available:\n",
    "                if bias_type == \"None\":\n",
    "                    b_layer = Variable(torch.zeros(self.output_size).to(self.device), requires_grad = False)\n",
    "                elif bias_type == \"constant\":\n",
    "                    b_layer = self.b_layer_seed[0].repeat(self.output_size)\n",
    "                elif bias_type == \"arithmetic-series\":\n",
    "                    mean = self.b_layer_seed.mean()\n",
    "                    offset = self.output_size / float(2) - 0.5\n",
    "                    idx = Variable(torch.FloatTensor(range(self.output_size)).to(self.device), requires_grad = False)\n",
    "                    b_layer = mean + self.b_interval * (idx - offset)\n",
    "                elif bias_type == \"arithmetic-series-2D\":\n",
    "                    idx_i, idx_j = np.meshgrid(range(self.output_size_2D[0]), range(self.output_size_2D[1]), indexing = \"ij\")\n",
    "                    idx_i = torch.from_numpy(idx_i).float().view(-1).to(self.device)\n",
    "                    idx_j = torch.from_numpy(idx_j).float().view(-1).to(self.device)\n",
    "                    offset_i = self.output_size_2D[0] / float(2) - 0.5\n",
    "                    offset_j = self.output_size_2D[1] / float(2) - 0.5\n",
    "                    b_layer = (self.b_mean_2D + \\\n",
    "                               self.b_interval_2D[0] * Variable(idx_i - offset_i, requires_grad = False) + \\\n",
    "                               self.b_interval_2D[1] * Variable(idx_j - offset_j, requires_grad = False)).view(-1)\n",
    "                elif bias_type == \"dense\":\n",
    "                    b_layer = self.b_layer_seed\n",
    "                else:\n",
    "                    raise Exception(\"bias_type '{0}' not recognized!\".format(bias_type))\n",
    "                self.b_list.append(b_layer)\n",
    "\n",
    "            if len(self.b_available) == 1:\n",
    "                self.b_core = b_layer\n",
    "            else:\n",
    "                self.b_list = torch.stack(self.b_list, dim = 1)\n",
    "                b_sig_softmax = nn.Softmax(dim = -1)(self.b_sig.unsqueeze(0))\n",
    "                self.b_core = torch.matmul(self.b_list, b_sig_softmax.transpose(1,0)).squeeze(1)\n",
    "\n",
    "\n",
    "    def forward(self, X, p_dict=None):\n",
    "        del p_dict\n",
    "        output = X\n",
    "        if hasattr(self, \"input_size_original\"):\n",
    "            output = output.view(-1, self.input_size)\n",
    "        # Get superposition of layers:\n",
    "        self.get_layers(source = [\"weight\", \"bias\"])\n",
    "\n",
    "        # Perform dot(X, W) + b:\n",
    "        output = torch.matmul(output, self.W_core) + self.b_core\n",
    "        \n",
    "        # Exert superposition of activation functions:\n",
    "        if len(self.A_available) == 1:\n",
    "            output = get_activation(self.A_available[0])(output)\n",
    "        else:\n",
    "            self.A_list = []\n",
    "            A_sig_softmax = nn.Softmax(dim = -1)(self.A_sig.unsqueeze(0))\n",
    "            for i, activation in enumerate(self.A_available):\n",
    "                A = get_activation(activation)(output)\n",
    "                self.A_list.append(A)\n",
    "            self.A_list = torch.stack(self.A_list, 2)\n",
    "            output = torch.matmul(self.A_list, A_sig_softmax.transpose(1,0)).squeeze(2)\n",
    "\n",
    "        if hasattr(self, \"output_size_original\"):\n",
    "            output = output.view(*((-1,) + self.output_size_original))\n",
    "        return output\n",
    "    \n",
    "    \n",
    "    def get_param_names(self, source):\n",
    "        if source == \"modules\":\n",
    "            param_names = [\"W_layer_seed\", \"b_layer_seed\"]\n",
    "            if \"arithmetic-series-in\" in self.W_available:\n",
    "                param_names.append(\"W_interval_j\")\n",
    "            if \"arithmetic-series-out\" in self.W_available:\n",
    "                param_names.append(\"W_interval_i\")\n",
    "            if \"arithmetic-series-2D-in\" in self.W_available:\n",
    "                param_names = param_names + [\"W_mean_2D_in\", \"W_interval_2D_in\"]\n",
    "            if \"arithmetic-series-2D-out\" in self.W_available:\n",
    "                param_names = param_names + [\"W_mean_2D_out\", \"W_interval_2D_out\"]\n",
    "            if \"arithmetic-series\" in self.b_available:\n",
    "                param_names.append(\"b_interval\")\n",
    "        if source == \"attention\":\n",
    "            param_names = [\"W_sig\", \"b_sig\", \"A_sig\"]\n",
    "        return param_names\n",
    "    \n",
    "    \n",
    "    def get_weights_bias(self):\n",
    "        self.get_layers(source = [\"weight\", \"bias\"])\n",
    "        return to_np_array(self.W_layer_seed), to_np_array(self.b_layer_seed)\n",
    "\n",
    "\n",
    "    def get_regularization(self, mode, source = [\"weight\", \"bias\"]):\n",
    "        reg = Variable(torch.FloatTensor(np.array([0])), requires_grad = False).to(self.device)\n",
    "        if not isinstance(source, list):\n",
    "            source = [source]\n",
    "        if mode == \"L1\":\n",
    "            if \"weight\" in source:\n",
    "                reg = reg + self.W_core.abs().sum()\n",
    "            if \"bias\" in source:\n",
    "                reg = reg + self.b_core.abs().sum()\n",
    "        elif mode == \"layer_L1\":\n",
    "            if \"weight\" in source:\n",
    "                self.get_layers(source = [\"weight\"])\n",
    "                reg = reg + self.W_list.abs().sum()\n",
    "            if \"bias\" in source:\n",
    "                self.get_layers(source = [\"bias\"])\n",
    "                reg = reg + self.b_list.abs().sum()\n",
    "        elif mode == \"L2\":\n",
    "            if \"weight\" in source:\n",
    "                reg = reg + torch.sum(self.W_core ** 2)\n",
    "            if \"bias\" in source:\n",
    "                reg = reg + torch.sum(self.b_core ** 2)\n",
    "        elif mode == \"S_entropy\":\n",
    "            if \"weight\" in source:\n",
    "                W_sig_softmax = nn.Softmax(dim = -1)(self.W_sig.unsqueeze(0))\n",
    "                reg = reg - torch.sum(W_sig_softmax * torch.log(W_sig_softmax))\n",
    "            if \"bias\" in source:\n",
    "                b_sig_softmax = nn.Softmax(dim = -1)(self.b_sig.unsqueeze(0))\n",
    "                reg = reg - torch.sum(b_sig_softmax * torch.log(b_sig_softmax))\n",
    "        elif mode == \"S_entropy_activation\":\n",
    "            A_sig_softmax = nn.Softmax(dim = -1)(self.A_sig.unsqueeze(0))\n",
    "            reg = reg - torch.sum(A_sig_softmax * torch.log(A_sig_softmax))\n",
    "        elif mode in AVAILABLE_REG:\n",
    "            pass\n",
    "        else:\n",
    "            raise Exception(\"mode '{}' not recognized!\".format(mode))\n",
    "        return reg\n",
    "\n",
    "\n",
    "    def set_cuda(self, is_cuda):\n",
    "        if isinstance(is_cuda, str):\n",
    "            self.cuda(is_cuda)\n",
    "        else:\n",
    "            if is_cuda:\n",
    "                self.cuda()\n",
    "            else:\n",
    "                self.cpu()\n",
    "        self.is_cuda = is_cuda"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}