{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3804e14d-4264-4ab9-97a0-7ff40839e70d",
   "metadata": {},
   "source": [
    "# 1. Design a layer that takes an input and computes a tensor reduction, i.e., it returns $y_k=\\sum_{i,j}W_{i,j,k}x_ix_j$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "b9cd241c-b2e1-44ef-b35a-98b52f7f4454",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torch\n",
    "\n",
    "class ReductionLayer(nn.Module):\n",
    "    def __init__(self, num_inputs, k):\n",
    "        super().__init__()\n",
    "        self.w = nn.ParameterList([nn.Parameter(torch.randn(num_inputs, num_inputs)) for i in range(k)])\n",
    "        \n",
    "    def stat_row(self, X):\n",
    "        y = []\n",
    "        for part_w in self.w:\n",
    "            y.append((part_w*X).sum().reshape(-1,1))\n",
    "        return torch.cat(y,dim=-1)\n",
    "        \n",
    "    def forward(self, X):\n",
    "        chunks = torch.chunk(X,X.shape[0],dim=0)\n",
    "        rows = []\n",
    "        for row in chunks:\n",
    "            row = row.reshape(1,-1)\n",
    "            part_x = torch.matmul(row.T,row)\n",
    "            rows.append(self.stat_row(part_x))\n",
    "        return torch.cat(rows,dim=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "04d4380b-62cd-4308-a452-6897c80e97be",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-4.3252, -3.5257],\n",
       "        [-5.1311, -0.4626]], grad_fn=<CatBackward0>)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "layer = ReductionLayer(5,2)\n",
    "x = torch.randn(2,5)\n",
    "layer(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e7afd68-dc81-48e2-b2d7-1ca7086311bc",
   "metadata": {},
   "source": [
    "# 2. Design a layer that returns the leading half of the Fourier coefficients of the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "976e4d36-d26e-4069-a33d-60757a9a39da",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input shape: torch.Size([1, 10, 10])\n",
      "Output shape: torch.Size([1, 10, 5])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "class FourierCoefficientsLayer(nn.Module):\n",
    "    def __init__(self, num_coefficients):\n",
    "        super(FourierCoefficientsLayer, self).__init__()\n",
    "        self.num_coefficients = num_coefficients\n",
    "    \n",
    "    def forward(self, x):\n",
    "        # Apply Fourier transform along the last dimension (assumed to be time dimension)\n",
    "        fourier_transform = torch.fft.fft(x)\n",
    "        \n",
    "        # Select the leading half of the coefficients\n",
    "        leading_coefficients = fourier_transform[..., :self.num_coefficients]\n",
    "        \n",
    "        return leading_coefficients\n",
    "\n",
    "# Create Fourier coefficients layer with 5 coefficients\n",
    "num_coefficients = 5\n",
    "fourier_layer = FourierCoefficientsLayer(num_coefficients)\n",
    "\n",
    "# Create example input with time dimension (e.g., audio signal)\n",
    "input_data = torch.randn(1, 10, 10)  # Batch size 1, 10 time steps, 2 features\n",
    "\n",
    "# Apply the Fourier coefficients layer\n",
    "output = fourier_layer(input_data)\n",
    "\n",
    "print(\"Input shape:\", input_data.shape)\n",
    "print(\"Output shape:\", output.shape)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:d2l]",
   "language": "python",
   "name": "conda-env-d2l-py"
  },
  "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}