From 98ad6ca1174a8fd78e53c1d99769b7fbc8017d99 Mon Sep 17 00:00:00 2001
From: Winter M <winter@antithesis.com>
Date: Mon, 6 Oct 2025 15:22:34 -0400
Subject: [PATCH] python3Packages.tensorflow-bin: fix cuda routines that
 require jit

Without this change (or $XLA_FLAGS set properly), errors like this would
crop up, resulting in JIT compilation failures:

2025-10-07 17:29:56.600986: W external/local_xla/xla/service/gpu/llvm_gpu_backend/default/nvptx_libdevice_path.cc:40] Can't find libdevice directory ${CUDA_DIR}/nvvm/libdevice. This may result in compilation or runtime failures, if the program we try to run uses routines from libdevice.
---
 pkgs/development/python-modules/tensorflow/bin.nix | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/pkgs/development/python-modules/tensorflow/bin.nix b/pkgs/development/python-modules/tensorflow/bin.nix
index ddbdedd266e1..f15160a2c31c 100644
--- a/pkgs/development/python-modules/tensorflow/bin.nix
+++ b/pkgs/development/python-modules/tensorflow/bin.nix
@@ -205,7 +205,12 @@ buildPythonPackage rec {
           ''}
         done
       ''
-    );
+    )
+    # Symlink nvcc besides TensorFlow so that routines that require JIT can work
+    # properly.
+    + lib.optionalString cudaSupport ''
+      ln -s ${cudaPackages.cuda_nvcc} "$out/${python.sitePackages}/tensorflow/cuda"
+    '';
 
   # Upstream has a pip hack that results in bin/tensorboard being in both tensorflow
   # and the propagated input tensorboard, which causes environment collisions.