feat: added the lstm demo so that it's tested as well through the dem…

…os tests (#73)

.github/workflows/_demo-test.yml

@@ -13,7 +13,15 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        modules: [alexnet_demo, bert_demo, image_segmentation_with_ivy_unet, mmpretrain_to_jax, resnet_demo, torch_to_jax]
+        modules: [
+          alexnet_demo,
+          bert_demo,
+          image_segmentation_with_ivy_unet,
+          mmpretrain_to_jax,
+          resnet_demo,
+          torch_to_jax,
+          lstm_demo,
+        ]
 
     steps:
         - name: Clean repository

examples_and_demos/lstm_demo.ipynb

@@ -0,0 +1,183 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
+      "\u001b[0m"
+     ]
+    }
+   ],
+   "source": [
+    "!pip install -q ivy"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2024-01-22 11:50:15.099601: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
+      "2024-01-22 11:50:15.099643: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
+      "2024-01-22 11:50:15.100960: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
+      "2024-01-22 11:50:15.108768: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
+      "To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
+      "2024-01-22 11:50:16.332076: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
+     ]
+    }
+   ],
+   "source": [
+    "import ivy\n",
+    "import torch\n",
+    "import tensorflow as tf\n",
+    "import numpy as np"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/workspaces/tracer-transpiler/ivy_repo/ivy/utils/exceptions.py:383: UserWarning: The current backend: 'tensorflow' does not support inplace updates natively. Ivy would quietly create new arrays when using inplace updates with this backend, leading to memory overhead (same applies for views). If you want to control your memory management, consider doing ivy.set_inplace_mode('strict') which should raise an error whenever an inplace update is attempted with this backend.\n",
+      "  warnings.warn(\n",
+      "2024-01-22 11:50:26.221321: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1929] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 14699 MB memory:  -> device: 0, name: Tesla V100-PCIE-16GB, pci bus id: 0001:00:00.0, compute capability: 7.0\n",
+      "2024-01-22 11:50:26.625856: I external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:454] Loaded cuDNN version 8902\n"
+     ]
+    }
+   ],
+   "source": [
+    "# create torch lstm layer\n",
+    "torch_lstm = torch.nn.LSTM(2, 2, 1).to(\"cuda\")\n",
+    "\n",
+    "# transpile lstm layer to tensorflow\n",
+    "x = torch.rand((5, 2, 2)).cuda()\n",
+    "tf_lstm = ivy.transpile(torch_lstm, source=\"torch\", to=\"tensorflow\", args=(x,))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "True"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# get output of original torch lstm layer\n",
+    "x = torch.rand((20, 32, 2)).cuda()\n",
+    "original_output = torch_lstm(x)\n",
+    "\n",
+    "# get output of transpiled tf lstm layer with the same input\n",
+    "x = tf.constant(x.cpu().numpy())\n",
+    "transpiled_output = tf_lstm(x)\n",
+    "\n",
+    "# verify the outputs are the same (with some tolerance)\n",
+    "np.allclose(original_output[0].detach().cpu(), transpiled_output[0].numpy(), atol=1e-7)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2024-01-22 11:50:30.698440: I external/local_tsl/tsl/platform/default/subprocess.cc:304] Start cannot spawn child process: No such file or directory\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "transpiled tf time is 3.725307350738295x slower than torch's lstm\n",
+      "original tf lstm time is 4.440824652724787x slower than torch's lstm\n"
+     ]
+    }
+   ],
+   "source": [
+    "# run some benchmarks\n",
+    "from time import perf_counter\n",
+    "\n",
+    "x = torch.rand((20, 32, 2)).cuda()\n",
+    "N_RUNS = 1000\n",
+    "\n",
+    "# time the original torch lstm\n",
+    "s = perf_counter()\n",
+    "for _ in range(N_RUNS):\n",
+    "  torch_lstm(x)\n",
+    "original_torch_time = perf_counter() - s\n",
+    "\n",
+    "# compile transpiled tf lstm\n",
+    "x = tf.constant(x.cpu().numpy())\n",
+    "tf_lstm = tf.autograph.experimental.do_not_convert(tf_lstm)\n",
+    "compiled_tf_lstm = tf.function(tf_lstm)\n",
+    "compiled_tf_lstm(x)\n",
+    "\n",
+    "# time the transpiled tf lstm\n",
+    "s = perf_counter()\n",
+    "for _ in range(N_RUNS):\n",
+    "  compiled_tf_lstm(x)\n",
+    "transpiled_tf_time = perf_counter() - s\n",
+    "\n",
+    "# time tf's own lstm layer (also compiled) for comparison\n",
+    "original_tf_lstm = tf.keras.layers.LSTM(2, time_major=True, return_sequences=True)\n",
+    "original_tf_lstm = tf.function(original_tf_lstm)\n",
+    "original_tf_lstm(x)\n",
+    "\n",
+    "s = perf_counter()\n",
+    "for _ in range(N_RUNS):\n",
+    "  original_tf_lstm(x)\n",
+    "original_tf_time = perf_counter() - s\n",
+    "\n",
+    "# as we can see, the transpiled tf lstm has comparable performance to tf's own lstm layer\n",
+    "print(f'transpiled tf time is {transpiled_tf_time / original_torch_time}x slower than torch\\'s lstm')\n",
+    "print(f'original tf lstm time is {original_tf_time / original_torch_time}x slower than torch\\'s lstm')"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "multienv",
+   "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.10.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}