trees forward pass done

Author: Ghost---Shadow

notebooks/library/stacks.py

@@ -43,3 +43,11 @@ def stack_pop(state):
     element = superposition_lookup_vectored(buffer, index)
     state = (buffer, index)
     return state, element
+
+
+@tf.function
+def stack_peek(stack):
+    buffer, index = stack
+    index = tf.roll(index, shift=-1, axis=0)
+    element = superposition_lookup_vectored(buffer, index)
+    return element

notebooks/stacks.ipynb

@@ -30,7 +30,9 @@
    "source": [
     "## Stack representation\n",
     "\n",
-    "The `stack` variable has two variables, buffer and index. The buffer is the writable buffer where stack elements are stored. Index points to top of stack + 1."
+    "The `stack` variable has two variables, buffer and index. The buffer is the writable buffer where stack elements are stored. Index points to top of stack + 1.\n",
+    "\n",
+    "Note: Since these functions can create variables, they must execute eagerly."
    ]
   },
   {
@@ -248,6 +250,51 @@
     "print(tape.gradient(element, buffer))"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Stack peek\n",
+    "\n",
+    "Get the stack top without any modification"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "tf.Tensor([3. 3. 3.], shape=(3,), dtype=float32)\n",
+      "None\n"
+     ]
+    }
+   ],
+   "source": [
+    "@tf.function\n",
+    "def stack_peek(stack):\n",
+    "    buffer, index = stack\n",
+    "    index = tf.roll(index, shift=-1, axis=0)\n",
+    "    element = superposition_lookup_vectored(buffer, index)\n",
+    "    return element\n",
+    "\n",
+    "buffer = tf.Variable([\n",
+    "    [1,1,1],\n",
+    "    [2,2,2],\n",
+    "    [3,3,3]\n",
+    "],dtype=tf.float32)\n",
+    "stack = new_stack_from_buffer(buffer, True)\n",
+    "\n",
+    "with tf.GradientTape() as tape:\n",
+    "    element = stack_peek(stack)\n",
+    "\n",
+    "print(element)\n",
+    "print(tape.gradient(element, buffer))"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -262,7 +309,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [
     {
@@ -328,7 +375,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [
     {