diff --git a/pkgs/development/python-modules/einops/default.nix b/pkgs/development/python-modules/einops/default.nix
index 18b8d671c98a..6678ff63efc5 100644
--- a/pkgs/development/python-modules/einops/default.nix
+++ b/pkgs/development/python-modules/einops/default.nix
@@ -7,6 +7,7 @@
 , nbconvert
 , jupyter
 , chainer
+, parameterized
 , pytorch
 , mxnet
 , tensorflow
@@ -15,13 +16,13 @@
 
 buildPythonPackage rec {
   pname = "einops";
-  version = "0.3.2";
+  version = "0.4.1";
 
   src = fetchFromGitHub {
     owner = "arogozhnikov";
     repo = pname;
     rev = "v${version}";
-    sha256 = "0ix094cfh6w4bvx6ymp5dpm35y9nkaibcn1y50g6kwdp4f0473y8";
+    hash = "sha256-n4R4lcRimuOncisCTs2zJWPlqZ+W2yPkvkWAnx4R91s=";
   };
 
   checkInputs = [
@@ -33,6 +34,7 @@ buildPythonPackage rec {
     jupyter
     # For backend tests
     chainer
+    parameterized
     pytorch
     mxnet
     tensorflow
diff --git a/pkgs/development/python-modules/elegy/default.nix b/pkgs/development/python-modules/elegy/default.nix
index 5b0cb293c0a3..a4d17a830e4d 100644
--- a/pkgs/development/python-modules/elegy/default.nix
+++ b/pkgs/development/python-modules/elegy/default.nix
@@ -22,16 +22,24 @@
 
 buildPythonPackage rec {
   pname = "elegy";
-  version = "0.8.4";
+  version = "0.8.6";
   format = "pyproject";
 
   src = fetchFromGitHub {
     owner = "poets-ai";
     repo = pname;
     rev = version;
-    sha256 = "11w8lgl31b52w2qri8j8cgzd30sn8i3769g8nkkshvgkjgca9r4g";
+    hash = "sha256-FZmLriYhsX+zyQKCtCjbOy6MH+AvjzHRNUyaDSXGlLI=";
   };
 
+  # The cloudpickle constraint is too strict. wandb is marked as an optional
+  # dependency but `buildPythonPackage` doesn't seem to respect that setting.
+  postPatch = ''
+    substituteInPlace pyproject.toml \
+      --replace 'cloudpickle = "^1.5.0"' 'cloudpickle = "*"' \
+      --replace 'wandb = { version = "^0.12.10", optional = true }' ""
+  '';
+
   nativeBuildInputs = [
     poetry
   ];
diff --git a/pkgs/development/python-modules/optax/default.nix b/pkgs/development/python-modules/optax/default.nix
index ae8cfdbbc282..f3790a3f59f4 100644
--- a/pkgs/development/python-modules/optax/default.nix
+++ b/pkgs/development/python-modules/optax/default.nix
@@ -8,19 +8,19 @@
 , numpy
 , pytest-xdist
 , pytestCheckHook
+, tensorflow
+, tensorflow-datasets
 }:
 
 buildPythonPackage rec {
   pname = "optax";
-  # As of 2022-01-06, the latest stable version (0.1.0) has broken tests that are fixed
-  # in https://github.com/deepmind/optax/commit/d6633365d84eb6f2c0df0c52b630481a349ce562
-  version = "unstable-2022-01-05";
+  version = "0.1.1";
 
   src = fetchFromGitHub {
     owner = "deepmind";
     repo = pname;
-    rev = "5ec5541b3486224b22e950480ff639ceaf5098f7";
-    sha256 = "1q8cxc42a5xais2ll1l238cnn3l7w28savhgiz0lg01ilz2ysbli";
+    rev = "v${version}";
+    hash = "sha256-s/BcqzhdfWzR61MStusUPQtuT4+t8NcC5gBGiGggFqw=";
   };
 
   buildInputs = [ jaxlib ];
@@ -35,6 +35,8 @@ buildPythonPackage rec {
     dm-haiku
     pytest-xdist
     pytestCheckHook
+    tensorflow
+    tensorflow-datasets
   ];
   pytestFlagsArray = [ "-n $NIX_BUILD_CORES" ];
 
@@ -45,8 +47,7 @@ buildPythonPackage rec {
   disabledTestPaths = [
     # Requires `flax` which depends on `optax` creating circular dependency.
     "optax/_src/equivalence_test.py"
-    # Require `tensorflow_datasets` which isn't packaged in `nixpkgs`.
-    "examples/datasets_test.py"
+    # See https://github.com/deepmind/optax/issues/323.
     "examples/lookahead_mnist_test.py"
   ];
 
diff --git a/pkgs/development/python-modules/treeo/default.nix b/pkgs/development/python-modules/treeo/default.nix
index f3271b9a3bb2..febb59880b69 100644
--- a/pkgs/development/python-modules/treeo/default.nix
+++ b/pkgs/development/python-modules/treeo/default.nix
@@ -1,5 +1,6 @@
 { buildPythonPackage
 , fetchFromGitHub
+, fetchpatch
 , jax
 , jaxlib
 , lib
@@ -8,20 +9,30 @@
 
 buildPythonPackage rec {
   pname = "treeo";
-  version = "0.4.0";
+  # Note that there is a version 0.4.0, but it was released in error. At the
+  # time of writing (2022-03-29), v0.0.11 is the latest as reported on GitHub
+  # and PyPI.
+  version = "0.0.11";
   format = "pyproject";
 
   src = fetchFromGitHub {
     owner = "cgarciae";
     repo = pname;
     rev = version;
-    sha256 = "176r1kgsdlylvdrxmhnzni81p8m9cfnsn4wwn6fnmsgam2qbp76j";
+    hash = "sha256-zs3F8i+G5OX/A9wOO60xVuvnm2QqrL+dHIrC0qwH37o=";
   };
 
-  postPatch = ''
-    substituteInPlace pyproject.toml \
-      --replace 'typing-extensions = "^3.10.0"' 'typing-extensions = "*"'
-  '';
+  # See https://github.com/cgarciae/treex/issues/68.
+  patches = [
+    (fetchpatch {
+      url = "https://github.com/cgarciae/treeo/pull/14/commits/022915da2b3bf76406a7c79d1b4593bee7956f16.patch";
+      hash = "sha256-WGxJqqrf2g0yZe30RyG1xxbloiqj1awuf1Y4eh5y+z0=";
+    })
+    (fetchpatch {
+      url = "https://github.com/cgarciae/treeo/pull/14/commits/99f9488bd0c977780844fd79743167b0010d359b.patch";
+      hash = "sha256-oKDYs+Ah0QXkhiJysIudQ6VLIiUiIcnQisxYp6GJuTc=";
+    })
+  ];
 
   nativeBuildInputs = [
     poetry-core
diff --git a/pkgs/development/python-modules/treex/default.nix b/pkgs/development/python-modules/treex/default.nix
index 7ed83adc64d1..fb0a0c7d87c7 100644
--- a/pkgs/development/python-modules/treex/default.nix
+++ b/pkgs/development/python-modules/treex/default.nix
@@ -18,21 +18,23 @@
 
 buildPythonPackage rec {
   pname = "treex";
-  version = "0.6.9";
+  version = "0.6.10";
   format = "pyproject";
 
   src = fetchFromGitHub {
     owner = "cgarciae";
     repo = pname;
     rev = version;
-    sha256 = "1yvlldmhji12h249j14ba44hnb9x1fhrj7rh1cx2vn0vxj5wpg7x";
+    hash = "sha256-ZHfgmRNbFh8DFZkmilY0pmRNQhJFqT689I7Lu8FuFm4=";
   };
 
+  # At the time of writing (2022-03-29), rich is currently at version 11.0.0.
+  # The treeo dependency is compatible with a patch, but not marked as such in
+  # treex. See https://github.com/cgarciae/treex/issues/68.
   postPatch = ''
     substituteInPlace pyproject.toml \
-      --replace 'rich = "^10.7.0"' 'rich = ">=10.7.0"' \
-      --replace 'PyYAML = "^5.4.1"' 'PyYAML = ">=5.4.1"' \
-      --replace 'optax = "^0.0.9"' 'optax = ">=0.0.9"'
+      --replace 'rich = "^11.2.0"' 'rich = "*"' \
+      --replace 'treeo = "^0.0.10"' 'treeo = "*"'
   '';
 
   nativeBuildInputs = [