From 2308a66c28be018307d2f0b8ddc10d9bf917b7aa Mon Sep 17 00:00:00 2001
From: Guillaume Yziquel <142336927+gl-yziquel@users.noreply.github.com>
Date: Fri, 21 Jun 2024 11:58:37 +0000
Subject: [PATCH] Fixing paths for python wheel generation.

---
 pip_pkg.sh | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pip_pkg.sh b/pip_pkg.sh
index e19dbbf4cd..d3d4c72de5 100755
--- a/pip_pkg.sh
+++ b/pip_pkg.sh
@@ -36,8 +36,8 @@ DEST=$(cd "$DEST" && pwd)
 # Use cp -L to resolve symlinks added as part of bazel build. Otherwise wheel
 # does not follow them and fails to include TFP on JAX/Numpy.
 pip_tmp="$(mktemp -d)"
-cp -L -R bazel-bin/pip_pkg.runfiles/tensorflow_probability "${pip_tmp}"
-cd "${pip_tmp}/tensorflow_probability"
+cp -L -R bazel-bin/pip_pkg.runfiles/_main "${pip_tmp}"
+cd "${pip_tmp}/_main"
 
 echo >>tensorflow_probability/python/version.py \
   "if __version__.endswith('dev'): __version__ += '$(date --utc +%Y%m%d)'"