chore: fix imports in notebooks

michellemli · Dec 11, 2024 · 1a0a2f4 · 1a0a2f4
1 parent 6f0dea2
commit 1a0a2f4
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 452 deletions.
diff --git a/notebooks/xlstm_large/demo.ipynb b/notebooks/xlstm_large/demo.ipynb
@@ -20,7 +20,7 @@
    "outputs": [],
    "source": [
     "from xlstm.xlstm_large.model import xLSTMLargeConfig, xLSTMLarge\n",
-    "from xlstm_kernels.torch import get_available_xlstm_step_kernels, get_available_xlstm_kernels, get_available_xlstm_sequence_kernels"
+    "from mlstm_kernels.torch import get_available_mlstm_step_kernels, get_available_mlstm_kernels, get_available_mlstm_sequence_kernels"
    ]
   },
   {
@@ -29,7 +29,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "get_available_xlstm_kernels(), get_available_xlstm_step_kernels(), get_available_xlstm_sequence_kernels()"
+    "get_available_mlstm_kernels(), get_available_mlstm_step_kernels(), get_available_mlstm_sequence_kernels()"
    ]
   },
   {
@@ -46,8 +46,8 @@
     "    return_last_states=True,\n",
     "    mode=\"inference\",\n",
     "    chunkwise_kernel=\"chunkwise--triton_limit_chunk\",\n",
-    "    sequence_kernel=\"native_sequence__triton_step_fused\",\n",
-    "    step_kernel=\"triton_fused\",\n",
+    "    sequence_kernel=\"native_sequence__triton\",\n",
+    "    step_kernel=\"triton\",\n",
     ")"
    ]
   },