diff --git a/pkgs/development/python-modules/optax/default.nix b/pkgs/development/python-modules/optax/default.nix
index b0f896a18ebb..42541d9078cf 100644
--- a/pkgs/development/python-modules/optax/default.nix
+++ b/pkgs/development/python-modules/optax/default.nix
@@ -1,28 +1,29 @@
 { absl-py
 , buildPythonPackage
 , chex
-, dm-haiku
 , fetchFromGitHub
 , jaxlib
 , lib
 , numpy
-, pytest-xdist
-, pytestCheckHook
-, tensorflow
-, tensorflow-datasets
+, callPackage
 }:
 
 buildPythonPackage rec {
   pname = "optax";
-  version = "0.1.1";
+  version = "0.1.3";
 
   src = fetchFromGitHub {
     owner = "deepmind";
     repo = pname;
     rev = "v${version}";
-    hash = "sha256-s/BcqzhdfWzR61MStusUPQtuT4+t8NcC5gBGiGggFqw=";
+    hash = "sha256-XAYztMBQpLBHNuNED/iodbwIMJSN/0GxdmTGQ5jD9Ws=";
   };
 
+  outputs = [
+    "out"
+    "testsout"
+  ];
+
   buildInputs = [ jaxlib ];
 
   propagatedBuildInputs = [
@@ -31,24 +32,21 @@ buildPythonPackage rec {
     numpy
   ];
 
-  checkInputs = [
-    dm-haiku
-    pytest-xdist
-    pytestCheckHook
-    tensorflow
-    tensorflow-datasets
-  ];
+  postInstall = ''
+    mkdir $testsout
+    cp -R examples $testsout/examples
+  '';
 
   pythonImportsCheck = [
     "optax"
   ];
 
-  disabledTestPaths = [
-    # Requires `flax` which depends on `optax` creating circular dependency.
-    "optax/_src/equivalence_test.py"
-    # See https://github.com/deepmind/optax/issues/323.
-    "examples/lookahead_mnist_test.py"
-  ];
+  # check in passthru.tests.pytest to escape infinite recursion with flax
+  doCheck = false;
+
+  passthru.tests = {
+    pytest = callPackage ./tests.nix { };
+  };
 
   meta = with lib; {
     description = "Optax is a gradient processing and optimization library for JAX.";
diff --git a/pkgs/development/python-modules/optax/tests.nix b/pkgs/development/python-modules/optax/tests.nix
new file mode 100644
index 000000000000..fedff61302d7
--- /dev/null
+++ b/pkgs/development/python-modules/optax/tests.nix
@@ -0,0 +1,35 @@
+{ stdenv
+, buildPythonPackage
+, dm-haiku
+, pytest-xdist
+, pytestCheckHook
+, tensorflow
+, tensorflow-datasets
+, flax
+, optax
+}:
+
+buildPythonPackage rec {
+  pname = "optax-tests";
+  inherit (optax) version;
+
+  src = optax.testsout;
+
+  dontBuild = true;
+  dontInstall = true;
+
+  checkInputs = [
+    dm-haiku
+    pytest-xdist
+    pytestCheckHook
+    tensorflow
+    tensorflow-datasets
+    flax
+  ];
+
+  disabledTestPaths = [
+    # See https://github.com/deepmind/optax/issues/323
+    "examples/lookahead_mnist_test.py"
+  ];
+
+}