JAX Tutorials: Fixing RNG JAX/PyTorch clash

Author: phlippe

docs/tutorial_notebooks/JAX/tutorial11/NF_image_modeling.ipynb

@@ -107,6 +107,8 @@
     "import jax\n",
     "import jax.numpy as jnp\n",
     "from jax import random\n",
+    "# Seeding for random operations\n",
+    "main_rng = random.PRNGKey(42)\n",
     "\n",
     "## Flax (NN in JAX)\n",
     "try:\n",
@@ -136,9 +138,6 @@
     "# Path to the folder where the pretrained models are saved\n",
     "CHECKPOINT_PATH = \"../../saved_models/tutorial11_jax\"\n",
     "\n",
-    "# Seeding for random operations\n",
-    "main_rng = random.PRNGKey(42)\n",
-    "\n",
     "print(\"Device:\", jax.devices()[0])"
    ]
   },

docs/tutorial_notebooks/JAX/tutorial12/Autoregressive_Image_Modeling.ipynb

@@ -104,6 +104,8 @@
     "import jax\n",
     "import jax.numpy as jnp\n",
     "from jax import random\n",
+    "# Seeding for random operations\n",
+    "main_rng = random.PRNGKey(42)\n",
     "\n",
     "## Flax (NN in JAX)\n",
     "try:\n",

docs/tutorial_notebooks/JAX/tutorial13/GAN.ipynb

@@ -107,6 +107,8 @@
     "import jax\n",
     "import jax.numpy as jnp\n",
     "from jax import random\n",
+    "# Seeding for random operations\n",
+    "main_rng = random.PRNGKey(42)\n",
     "\n",
     "## Flax (NN in JAX)\n",
     "try:\n",
@@ -140,9 +142,6 @@
     "# Path to the folder where the pretrained models are saved\n",
     "CHECKPOINT_PATH = \"../../saved_models/tutorial15_jax\"\n",
     "\n",
-    "# Seeding for random operations\n",
-    "main_rng = random.PRNGKey(42)\n",
-    "\n",
     "print(\"Device:\", jax.devices()[0])"
    ]
   },

docs/tutorial_notebooks/JAX/tutorial15/Vision_Transformer.ipynb

@@ -104,6 +104,8 @@
     "import jax\n",
     "import jax.numpy as jnp\n",
     "from jax import random\n",
+    "# Seeding for random operations\n",
+    "main_rng = random.PRNGKey(42)\n",
     "\n",
     "## Flax (NN in JAX)\n",
     "try:\n",

docs/tutorial_notebooks/JAX/tutorial17/SimCLR.ipynb

@@ -84,6 +84,8 @@
     "import jax.numpy as jnp\n",
     "from jax import random\n",
     "from jax.tree_util import tree_map\n",
+    "# Seeding for random operations\n",
+    "main_rng = random.PRNGKey(42)\n",
     "\n",
     "## Flax (NN in JAX)\n",
     "try:\n",
@@ -229,7 +231,7 @@
     "# We define a set of data loaders that we can use for various purposes later.\n",
     "# Note that for actually training a model, we will use different data loaders\n",
     "# with a lower batch size.\n",
-    "train_loader = data.DataLoader(train_set,\n",
+    "train_loader = data.DataLoader(train_set,r\n",
     "                               batch_size=1024,\n",
     "                               shuffle=False,\n",
     "                               drop_last=False,\n",

docs/tutorial_notebooks/JAX/tutorial4/Optimization_and_Initialization.ipynb

@@ -92,6 +92,8 @@
     "import jax\n",
     "import jax.numpy as jnp\n",
     "from jax import random\n",
+    "# Seeding for random operations\n",
+    "main_rng = random.PRNGKey(42)\n",
     "\n",
     "## Flax (NN in JAX)\n",
     "try:\n",
@@ -144,9 +146,6 @@
     "# Path to the folder where the pretrained models are saved\n",
     "CHECKPOINT_PATH = \"../../saved_models/tutorial5_jax\"\n",
     "\n",
-    "# Seeding for random operations\n",
-    "main_rng = random.PRNGKey(42)\n",
-    "\n",
     "print(\"Device:\", jax.devices()[0])"
    ]
   },

docs/tutorial_notebooks/JAX/tutorial5/Inception_ResNet_DenseNet.ipynb

@@ -101,6 +101,8 @@
     "import jax\n",
     "import jax.numpy as jnp\n",
     "from jax import random\n",
+    "# Seeding for random operations\n",
+    "main_rng = random.PRNGKey(42)\n",
     "\n",
     "## Flax (NN in JAX)\n",
     "try:\n",
@@ -131,9 +133,6 @@
     "# Path to the folder where the pretrained models are saved\n",
     "CHECKPOINT_PATH = \"../../saved_models/tutorial6_jax\"\n",
     "\n",
-    "# Seeding for random operations\n",
-    "main_rng = random.PRNGKey(42)\n",
-    "\n",
     "print(\"Device:\", jax.devices()[0])"
    ]
   },

docs/tutorial_notebooks/JAX/tutorial6/Transformers_and_MHAttention.ipynb

@@ -90,6 +90,8 @@
     "import jax\n",
     "import jax.numpy as jnp\n",
     "from jax import random\n",
+    "# Seeding for random operations\n",
+    "main_rng = random.PRNGKey(42)\n",
     "\n",
     "## Flax (NN in JAX)\n",
     "try:\n",
@@ -119,9 +121,6 @@
     "# Path to the folder where the pretrained models are saved\n",
     "CHECKPOINT_PATH = \"../../saved_models/tutorial7_jax\"\n",
     "\n",
-    "# Seeding for random operations\n",
-    "main_rng = random.PRNGKey(42)\n",
-    "\n",
     "print(\"Device:\", jax.devices()[0])"
    ]
   },
@@ -535,7 +534,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "Python 3",
    "language": "python",
    "name": "python3"
   },
@@ -549,7 +548,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.2"
+   "version": "3.7.9"
   }
  },
  "nbformat": 4,

docs/tutorial_notebooks/JAX/tutorial7/GNN_overview.ipynb

@@ -100,6 +100,8 @@
     "import jax\n",
     "import jax.numpy as jnp\n",
     "from jax import random\n",
+    "# Seeding for random operations\n",
+    "main_rng = random.PRNGKey(42)\n",
     "\n",
     "## Flax (NN in JAX)\n",
     "try:\n",