From 158b9d526d7adebfaedab83604ab961c824edb33 Mon Sep 17 00:00:00 2001
From: Samuel Ainsworth <skainsworth@gmail.com>
Date: Mon, 11 Apr 2022 20:57:28 +0000
Subject: [PATCH 1/2] python3Packages.chex: unstable-2021-12-16 -> 0.1.2

---
 pkgs/development/python-modules/chex/default.nix | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/pkgs/development/python-modules/chex/default.nix b/pkgs/development/python-modules/chex/default.nix
index bb873f2c3bfd..a06cb4ec84bb 100644
--- a/pkgs/development/python-modules/chex/default.nix
+++ b/pkgs/development/python-modules/chex/default.nix
@@ -12,16 +12,14 @@
 
 buildPythonPackage rec {
   pname = "chex";
-  # As of 2021-12-29, the latest official version has broken tests with jax 0.2.26:
-  # `AttributeError: module 'jax.interpreters.xla' has no attribute 'xb'`
-  version = "unstable-2021-12-16";
+  version = "0.1.2";
   format = "setuptools";
 
   src = fetchFromGitHub {
     owner = "deepmind";
     repo = pname;
-    rev = "5adc10e0b4218f8ec775567fca38b68bbad42a3a";
-    sha256 = "00xib6zv9pwid2q7wcr109qj3fa3g3b852skz8444kw7r0qxy7z3";
+    rev = "v${version}";
+    hash = "sha256-NtZYOHByKBcKmhRaNULwaQqxfoPRmgbtJ3cFHNfy4E8=";
   };
 
   propagatedBuildInputs = [

From 8442d63c5b662958c9ceb1bbf08bb492bf4a6d9b Mon Sep 17 00:00:00 2001
From: Samuel Ainsworth <skainsworth@gmail.com>
Date: Mon, 11 Apr 2022 21:07:50 +0000
Subject: [PATCH 2/2] python3Packages.dm-haiku: fix build

---
 pkgs/development/python-modules/dm-haiku/default.nix | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/pkgs/development/python-modules/dm-haiku/default.nix b/pkgs/development/python-modules/dm-haiku/default.nix
index 03677faa689f..8a6eabfde0f5 100644
--- a/pkgs/development/python-modules/dm-haiku/default.nix
+++ b/pkgs/development/python-modules/dm-haiku/default.nix
@@ -21,7 +21,7 @@ buildPythonPackage rec {
     owner = "deepmind";
     repo = pname;
     rev = "v${version}";
-    sha256 = "sha256-qvKMeGPiWXvvyV+GZdTWdsC6Wp08AmP8nDtWk7sZtqM=";
+    hash = "sha256-qvKMeGPiWXvvyV+GZdTWdsC6Wp08AmP8nDtWk7sZtqM=";
   };
 
   propagatedBuildInputs = [
@@ -32,6 +32,7 @@ buildPythonPackage rec {
   checkInputs = [
     chex
     cloudpickle
+    dill
     dm-tree
     jaxlib
     pytest-xdist
@@ -55,6 +56,11 @@ buildPythonPackage rec {
     "haiku/_src/integration/jax2tf_test.py"
   ];
 
+  disabledTests = [
+    # See https://github.com/deepmind/dm-haiku/issues/366.
+    "test_jit_Recurrent"
+  ];
+
   meta = with lib; {
     description = "Haiku is a simple neural network library for JAX developed by some of the authors of Sonnet.";
     homepage = "https://github.com/deepmind/dm-haiku";