From 672f7edb4d7a14e6d2e8026f49c66be270818b0a Mon Sep 17 00:00:00 2001
From: Samuel Ainsworth <skainsworth@gmail.com>
Date: Mon, 11 Apr 2022 19:51:45 +0000
Subject: [PATCH] python3Packages.tensorflow-datasets: fix build

---
 .../python-modules/tensorflow-datasets/default.nix         | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/pkgs/development/python-modules/tensorflow-datasets/default.nix b/pkgs/development/python-modules/tensorflow-datasets/default.nix
index 1fc9f5ae25dc..959ecd9092ab 100644
--- a/pkgs/development/python-modules/tensorflow-datasets/default.nix
+++ b/pkgs/development/python-modules/tensorflow-datasets/default.nix
@@ -9,6 +9,8 @@
 , future
 , imagemagick
 , importlib-resources
+, jax
+, jaxlib
 , jinja2
 , langdetect
 , lib
@@ -78,6 +80,8 @@ buildPythonPackage rec {
     beautifulsoup4
     ffmpeg
     imagemagick
+    jax
+    jaxlib
     jinja2
     langdetect
     matplotlib
@@ -119,6 +123,9 @@ buildPythonPackage rec {
     # Requires `tensorflow_io` which is not packaged in `nixpkgs`.
     "tensorflow_datasets/image/lsun_test.py"
 
+    # Requires `envlogger` which is not packaged in `nixpkgs`.
+    "tensorflow_datasets/rlds/robosuite_panda_pick_place_can/robosuite_panda_pick_place_can_test.py"
+
     # Fails with `TypeError: Constant constructor takes either 0 or 2 positional arguments`
     # deep in TF AutoGraph. Doesn't reproduce in Docker with Ubuntu 22.04 => might be related
     # to the differences in some of the dependencies?