pushing dirty work

jeertmans · May 31, 2024 · 8e9b318 · 8e9b318
1 parent 42aaec6
commit 8e9b318
Show file tree

Hide file tree

Showing 3 changed files with 1,445 additions and 221 deletions.
diff --git a/docs/source/GFlowNet.ipynb b/docs/source/GFlowNet.ipynb
diff --git a/docs/source/NotGFlowNet.ipynb b/docs/source/NotGFlowNet.ipynb
diff --git a/docs/source/Untitled.ipynb b/docs/source/Untitled.ipynb
@@ -0,0 +1,131 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "7ba6bab9-2eab-4493-a57a-c2583c41c953",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import jax\n",
+    "import jax.numpy as jnp"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "2e23ef34-c527-4bee-b957-28930158418b",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "Array(1, dtype=int32)"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "key = jax.random.PRNGKey(1234)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "b6b8555e-9cd6-4498-8041-e5cb3a74efd5",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(Array([1, 2, 3], dtype=int32), Array([5002, 2482, 2516], dtype=int32))"
+      ]
+     },
+     "execution_count": 19,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "key, key_cat = jax.random.split(key, 2)\n",
+    "logits = jnp.array([0.0, 2.0, 1.0, 1.0])\n",
+    "jnp.unique(jax.random.categorical(key_cat, jnp.log(logits), shape=(10000,)), return_counts=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "id": "346b2257-aa9c-4f54-aea1-422ec7bcfd6f",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(Array([1, 2, 3], dtype=int32), Array([5002, 2482, 2516], dtype=int32))"
+      ]
+     },
+     "execution_count": 21,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "jnp.unique(jax.random.categorical(key_cat, jnp.log(logits), shape=(10000,)), return_counts=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 28,
+   "id": "f8c51678-2404-43d7-93af-30362863a438",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "Array([[0, 1, 2, 3, 4],\n",
+       "       [0, 1, 2, 3, 4]], dtype=int32)"
+      ]
+     },
+     "execution_count": 28,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "jnp.tile(jnp.arange(5), (2, 1))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4382fa4a-04c2-4ae4-aace-c8728a4acd11",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "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.11.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}