FITS_onnx.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ONNX implementation of FITS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from onnx import TensorProto, numpy_helper\n",
    "from onnx.helper import make_node, make_graph, make_model, make_tensor_value_info\n",
    "from onnx.checker import check_model\n",
    "import numpy as np\n",
    "import torch\n",
    "from onnx.reference import ReferenceEvaluator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights_real=np.ones((2,4)).astype(np.float32)\n",
    "weights_imag=np.ones((2,4)).astype(np.float32)\n",
    "\n",
    "\n",
    "\n",
    "WR=numpy_helper.from_array(weights_real, name = 'WR')\n",
    "WI=numpy_helper.from_array(weights_imag, name = 'WI')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1., 1., 1., 1.],\n",
       "       [1., 1., 1., 1.]], dtype=float32)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "weights_imag"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## prepare variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = make_tensor_value_info('X', TensorProto.FLOAT, [None, None])\n",
    "# DSR = make_tensor_value_info('DSR', TensorProto.FLOAT, [None, None])\n",
    "# nomean= make_tensor_value_info('x_nomean', TensorProto.FLOAT, [None, None])\n",
    "# var= make_tensor_value_info('x_var', TensorProto.FLOAT, [None, None])\n",
    "x_normed = make_tensor_value_info('x_normed', TensorProto.FLOAT, [None, None])\n",
    "x_rfft = make_tensor_value_info('x_rfft', TensorProto.FLOAT, [None, None])\n",
    "x_rfft_imag = make_tensor_value_info('x_rfft_lpf_imag_perm', TensorProto.FLOAT, [None, None])\n",
    "x_rfft_real = make_tensor_value_info('x_rfft_lpf_real_perm', TensorProto.FLOAT, [None, None])\n",
    "\n",
    "out_R = make_tensor_value_info('R_perm', TensorProto.FLOAT, [None, None])\n",
    "out_I = make_tensor_value_info('I_perm', TensorProto.FLOAT, [None, None])\n",
    "RR = make_tensor_value_info('RR', TensorProto.FLOAT, [None, None])\n",
    "WWR = make_tensor_value_info('WR', TensorProto.FLOAT, [None, None])\n",
    "II = make_tensor_value_info('II', TensorProto.FLOAT, [None, None])\n",
    "WWI = make_tensor_value_info('WI', TensorProto.FLOAT, [None, None])\n",
    "\n",
    "out_rfft = make_tensor_value_info('out_rfft', TensorProto.FLOAT, [None, None])\n",
    "out = make_tensor_value_info('out', TensorProto.FLOAT, [None, None])\n",
    "\n",
    "\n",
    "L_axis = numpy_helper.from_array(np.array([1]), name='L_axis')\n",
    "LPF_start = numpy_helper.from_array(np.array([2]), name='LPF_start')\n",
    "LPF_end = numpy_helper.from_array(np.array([4]), name='LPF_end')\n",
    "# Perm_index = numpy_helper.from_array(np.array([0,2,1]), name='Perm_index')\n",
    "zero = numpy_helper.from_array(np.array([0]), name='zero')\n",
    "pad_Setting = numpy_helper.from_array(np.array([0,4,0,0,3,0]), name='pad_Setting')\n",
    "\n",
    "neg_one = numpy_helper.from_array(np.array([-1]), name='neg_one')\n",
    "neg_two = numpy_helper.from_array(np.array([-2]), name='neg_two')\n",
    "one = numpy_helper.from_array(np.array([1]), name='one')\n",
    "\n",
    "DSR = numpy_helper.from_array(np.array([3]).astype(np.float32), name='DSR')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the calculation graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# RIN\n",
    "# calculate the mean of the input tensor and remove the mean from the input with onnx\n",
    "Get_mean=make_node('ReduceMean', ['X', 'L_axis'], ['xmean'], keepdims=1)\n",
    "Get_nomean=make_node('Sub', ['X', 'xmean'], ['x_nomean']) # X_nomean = X - Xmean\n",
    "# calculate the sqrt(E(X-EX)^2) with onnx\n",
    "Get_nomean_square=make_node('Mul', ['x_nomean', 'x_nomean'], ['x_nomean_square'])\n",
    "Get_nomean_square_mean=make_node('ReduceMean', ['x_nomean_square', 'L_axis'], ['x_nomean_square_sum'], keepdims=1)\n",
    "Get_sqrt=make_node('Sqrt', ['x_nomean_square_sum'], ['x_var'])\n",
    "# calculate X_nomean/Xsqrt with onnx\n",
    "Get_x_normed=make_node('Div', ['x_nomean', 'x_var'], ['x_normed'])\n",
    "\n",
    "# LPF\n",
    "Get_rfft = make_node('DFT', ['x_normed'], ['x_rfft'], axis=1, onesided=1, inverse=0)\n",
    "Get_LPF = make_node('Slice', ['x_rfft', 'LPF_start', 'LPF_end', 'L_axis'], ['x_rfft_lpf'])\n",
    "Get_Split_Complex = make_node('Split', ['x_rfft_lpf'], ['x_rfft_lpf_real', 'x_rfft_lpf_imag'], axis=-1, num_outputs=2)\n",
    "\n",
    "# Complex Linear\n",
    "Get_PermuteI = make_node('Transpose', ['x_rfft_lpf_imag'], ['x_rfft_lpf_imag_perm'], perm=[0,2,1])\n",
    "Get_PermuteR = make_node('Transpose', ['x_rfft_lpf_real'], ['x_rfft_lpf_real_perm'], perm=[0,2,1])\n",
    "\n",
    "Get_RR = make_node('MatMul', ['x_rfft_lpf_real_perm', 'WR'], ['RR'])\n",
    "Get_RI = make_node('MatMul', ['x_rfft_lpf_real_perm', 'WI'], ['RI'])\n",
    "Get_IR = make_node('MatMul', ['x_rfft_lpf_imag_perm', 'WR'], ['IR'])\n",
    "Get_II = make_node('MatMul', ['x_rfft_lpf_imag_perm', 'WI'], ['II'])\n",
    "\n",
    "Get_R = make_node('Sub', ['RR', 'II'], ['R'])\n",
    "Get_I = make_node('Add', ['RI', 'IR'], ['I'])\n",
    "\n",
    "Get_inverse_PermuteI = make_node('Transpose', ['I'], ['I_perm'], perm=[0,2,1])\n",
    "Get_inverse_PermuteR = make_node('Transpose', ['R'], ['R_perm'], perm=[0,2,1])\n",
    "\n",
    "# Amplitude compensate\n",
    "Get_compensate_R = make_node('Mul', ['R_perm', 'DSR'], ['R_compensate'])\n",
    "Get_compensate_I = make_node('Mul', ['I_perm', 'DSR'], ['I_compensate'])\n",
    "\n",
    "# Zero padding\n",
    "Get_R_pad = make_node('Pad', ['R_compensate', 'pad_Setting', 'zero'], ['R_pad'], mode='constant')\n",
    "Get_I_pad = make_node('Pad', ['I_compensate', 'pad_Setting', 'zero'], ['I_pad'], mode='constant')\n",
    "\n",
    "# Prepare rfft\n",
    "Get_reverse_real = make_node('Slice', ['R_pad', 'neg_two', 'zero', 'one', 'neg_one'], ['R_rev'])\n",
    "Get_reverse_imag = make_node('Slice', ['I_pad', 'neg_two', 'zero', 'one', 'neg_one'], ['I_rev'])\n",
    "\n",
    "Get_neg = make_node('Neg', ['I_rev'], ['I_rev_neg']) # self conjugate\n",
    "Get_rev_Concat = make_node('Concat', ['R_rev', 'I_rev_neg'], ['x_rfft_rev'], axis=-1)\n",
    "Get_Concat = make_node('Concat', ['R_pad', 'I_pad'], ['x_rfft_forward'], axis=-1)\n",
    "\n",
    "Get_fin_Concat = make_node('Concat', ['x_rfft_forward', 'x_rfft_rev'], ['out_rfft'], axis=1) # 0-freq, forward spec, reverse spec\n",
    "\n",
    "# Inverse rfft, only accept full complex spectrum\n",
    "Get_inverse_rfft = make_node('DFT', ['out_rfft'], ['out_complex'], axis=1, inverse=1) \n",
    "Get_Temporal = make_node('Slice', ['out_complex', 'zero', 'one', 'neg_one'], ['out_temporal']) # only keep temporal domain\n",
    "\n",
    "# Rev RIN\n",
    "Get_out_nomean=make_node('Mul', ['out_temporal', 'x_var'], ['out_nomean']) \n",
    "Get_out=make_node('Add', ['out_nomean', 'xmean'], ['out'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make the graph (GraphProto)\n",
    "graph = make_graph(\n",
    "    [Get_mean, Get_nomean, Get_nomean_square, Get_nomean_square_mean, Get_sqrt, Get_x_normed, Get_rfft, Get_LPF, Get_Split_Complex, Get_PermuteI, Get_PermuteR,\n",
    "     Get_RR, Get_RI, Get_IR, Get_II, Get_R, Get_I, Get_inverse_PermuteI, Get_inverse_PermuteR, Get_compensate_I, Get_compensate_R ,Get_R_pad, Get_I_pad,\n",
    "     Get_reverse_real, Get_reverse_imag, Get_neg, Get_rev_Concat, Get_Concat, Get_fin_Concat, Get_inverse_rfft, Get_Temporal,\n",
    "     Get_out_nomean, Get_out],\n",
    "    'test-model',\n",
    "    [X],\n",
    "    [x_rfft, x_rfft_imag, x_rfft_real, out_R, out_I, out_rfft, out],\n",
    "    [L_axis,LPF_start,LPF_end,WR,WI, pad_Setting, zero, DSR, neg_one, neg_two, one]\n",
    ")\n",
    "\n",
    "onnx_model = make_model(graph, producer_name='onnx-example')\n",
    "\n",
    "check_model(onnx_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(2, 6, 2)\n"
     ]
    }
   ],
   "source": [
    "sees = ReferenceEvaluator(onnx_model)\n",
    "x=np.random.rand(2,10,1).astype(np.float32)\n",
    "xx=sees.run(None,{'X':x})\n",
    "print(xx[0].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Serialize and save the model as onnx file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the onnx model with example input\n",
    "with open(\"test.onnx\", \"wb\") as f:\n",
    "    f.write(onnx_model.SerializeToString())\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## check the input and output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[-2.113672  ,  3.2053206 ]],\n",
       "\n",
       "       [[-0.00395393,  3.129188  ]]], dtype=float32)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "realin=xx[2]\n",
    "realin # real part of the projection layer input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[-4.205752 ,  2.6030076]],\n",
       "\n",
       "       [[-2.4222276, -4.660607 ]]], dtype=float32)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "imagin=xx[1]\n",
    "imagin # imag part of the projection layer input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[0.3036429 ],\n",
       "        [0.60120684],\n",
       "        [0.48510543],\n",
       "        [0.75284505],\n",
       "        [0.23437363],\n",
       "        [0.26324   ],\n",
       "        [0.96830004],\n",
       "        [0.81264067],\n",
       "        [0.17516446],\n",
       "        [0.42030957]],\n",
       "\n",
       "       [[0.7057335 ],\n",
       "        [0.96225375],\n",
       "        [0.03172521],\n",
       "        [0.28581664],\n",
       "        [0.5284767 ],\n",
       "        [0.15667306],\n",
       "        [0.5521558 ],\n",
       "        [0.9609049 ],\n",
       "        [0.6574765 ],\n",
       "        [0.2306728 ]]], dtype=float32)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x # input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[ 1.337628  ],\n",
       "        [ 0.52596164],\n",
       "        [-0.14773268],\n",
       "        [ 0.6264461 ],\n",
       "        [ 0.6940575 ],\n",
       "        [ 0.5016829 ],\n",
       "        [ 0.6142318 ],\n",
       "        [ 0.32213527],\n",
       "        [ 0.39940047],\n",
       "        [ 0.6208314 ],\n",
       "        [ 0.5016829 ],\n",
       "        [ 0.62821245],\n",
       "        [ 0.47480467],\n",
       "        [ 0.28371495],\n",
       "        [ 0.5676295 ],\n",
       "        [ 0.5016829 ],\n",
       "        [ 0.64745516],\n",
       "        [ 0.77443516],\n",
       "        [-0.07232857],\n",
       "        [ 0.23172596]],\n",
       "\n",
       "       [[ 4.2864046 ],\n",
       "        [ 1.3112125 ],\n",
       "        [-2.606667  ],\n",
       "        [ 0.72204316],\n",
       "        [ 1.4868536 ],\n",
       "        [ 0.5071889 ],\n",
       "        [ 1.1259707 ],\n",
       "        [-0.21387142],\n",
       "        [-0.13313806],\n",
       "        [ 1.0284302 ],\n",
       "        [ 0.5071889 ],\n",
       "        [ 1.0966313 ],\n",
       "        [ 0.56359494],\n",
       "        [-0.5688746 ],\n",
       "        [ 0.695366  ],\n",
       "        [ 0.5071889 ],\n",
       "        [ 1.056249  ],\n",
       "        [ 2.0894585 ],\n",
       "        [-1.9099339 ],\n",
       "        [-1.4075183 ]]], dtype=float32)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xx[-1] # output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[-7.4505806e-07,  0.0000000e+00],\n",
       "        [-1.4925214e+00,  3.3711782e-01],\n",
       "        [-2.1136720e+00, -4.2057519e+00],\n",
       "        [ 3.2053206e+00,  2.6030076e+00],\n",
       "        [-2.1069021e+00, -7.1971887e-01],\n",
       "        [-2.6442459e+00,  0.0000000e+00]],\n",
       "\n",
       "       [[-1.0281801e-06,  0.0000000e+00],\n",
       "        [ 1.5157268e+00,  2.6613832e+00],\n",
       "        [-3.9539281e-03, -2.4222276e+00],\n",
       "        [ 3.1291881e+00, -4.6606069e+00],\n",
       "        [-1.2275178e+00, -1.2868351e+00],\n",
       "        [-3.9140186e-01,  4.4408921e-16]]], dtype=float32)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xx[0] # rfft of the input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[-0.51109576],\n",
       "        [-0.51109576],\n",
       "        [-0.51109576],\n",
       "        [-0.51109576]],\n",
       "\n",
       "       [[-3.9576    ],\n",
       "        [-3.9576    ],\n",
       "        [-3.9576    ],\n",
       "        [-3.9576    ]]], dtype=float32)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xx[-3] # RR RI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[ 1.337628  ],\n",
       "        [ 0.52596164],\n",
       "        [-0.14773268],\n",
       "        [ 0.6264461 ],\n",
       "        [ 0.6940575 ],\n",
       "        [ 0.5016829 ],\n",
       "        [ 0.6142318 ],\n",
       "        [ 0.32213527],\n",
       "        [ 0.39940047],\n",
       "        [ 0.6208314 ],\n",
       "        [ 0.5016829 ],\n",
       "        [ 0.62821245],\n",
       "        [ 0.47480467],\n",
       "        [ 0.28371495],\n",
       "        [ 0.5676295 ],\n",
       "        [ 0.5016829 ],\n",
       "        [ 0.64745516],\n",
       "        [ 0.77443516],\n",
       "        [-0.07232857],\n",
       "        [ 0.23172596]],\n",
       "\n",
       "       [[ 4.2864046 ],\n",
       "        [ 1.3112125 ],\n",
       "        [-2.606667  ],\n",
       "        [ 0.72204316],\n",
       "        [ 1.4868536 ],\n",
       "        [ 0.5071889 ],\n",
       "        [ 1.1259707 ],\n",
       "        [-0.21387142],\n",
       "        [-0.13313806],\n",
       "        [ 1.0284302 ],\n",
       "        [ 0.5071889 ],\n",
       "        [ 1.0966313 ],\n",
       "        [ 0.56359494],\n",
       "        [-0.5688746 ],\n",
       "        [ 0.695366  ],\n",
       "        [ 0.5071889 ],\n",
       "        [ 1.056249  ],\n",
       "        [ 2.0894585 ],\n",
       "        [-1.9099339 ],\n",
       "        [-1.4075183 ]]], dtype=float32)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xx[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[-0.51109576],\n",
       "        [-0.51109576],\n",
       "        [-0.51109576],\n",
       "        [-0.51109576]],\n",
       "\n",
       "       [[-3.9576    ],\n",
       "        [-3.9576    ],\n",
       "        [-3.9576    ],\n",
       "        [-3.9576    ]]], dtype=float32)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xx[4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "matmul: Input operand 1 has a mismatch in its core dimension 0, with gufunc signature (n?,k),(k,m?)->(n?,m?) (size 2 is different from 4)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_3642644/1365034337.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0moutr\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mxx\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtranspose\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0mouti\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mxx\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtranspose\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m# imag\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 4\u001b[0;31m \u001b[0moutR\u001b[0m\u001b[0;34m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmatmul\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutr\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mweights_real\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmatmul\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mouti\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mweights_imag\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m: matmul: Input operand 1 has a mismatch in its core dimension 0, with gufunc signature (n?,k),(k,m?)->(n?,m?) (size 2 is different from 4)"
     ]
    }
   ],
   "source": [
    "# permute the 1 and 2 axis of xx[1]\n",
    "outr=xx[3].transpose(0,2,1)\n",
    "outi=xx[2].transpose(0,2,1) # imag\n",
    "outR= np.matmul(outr,weights_real)-np.matmul(outi,weights_imag)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[-2.1290371, -2.1290371, -2.1290371, -2.1290371]],\n",
       "\n",
       "       [[-4.4575706, -4.4575706, -4.4575706, -4.4575706]]], dtype=float32)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "outR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[-4.0712576, -4.0712576, -4.0712576, -4.0712576]],\n",
       "\n",
       "       [[-2.6118875, -2.6118875, -2.6118875, -2.6118875]]], dtype=float32)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "outI = np.matmul(outr,weights_imag)+ np.matmul(outi,weights_real)\n",
    "outI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[ 8.5867941e-07,  0.0000000e+00],\n",
       "        [-5.2885100e-02,  4.9447557e-01],\n",
       "        [ 1.1048088e+00,  1.6531996e+00],\n",
       "        [-2.1126856e-01, -1.6675388e+00],\n",
       "        [-2.1087587e+00,  5.7152481e+00],\n",
       "        [-3.4242604e+00,  1.1102230e-16]],\n",
       "\n",
       "       [[ 8.5681677e-08,  0.0000000e+00],\n",
       "        [-3.2358274e+00, -9.8743612e-01],\n",
       "        [ 4.3768554e+00,  2.5047356e-01],\n",
       "        [ 3.4072566e+00, -2.0510821e+00],\n",
       "        [ 1.6312933e+00, -4.0278584e-01],\n",
       "        [ 1.1790700e+00,  2.2204460e-16]]], dtype=float32)"
      ]
     },
     "execution_count": 232,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xx[1]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "neuralforecast",
   "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.8.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}