From 8fde78e4c2bacf07b76a9aadb57e03b807c5e85d Mon Sep 17 00:00:00 2001
From: zjgarvey <zjgarvey@gmail.com>
Date: Fri, 3 Oct 2025 13:28:31 -0700
Subject: [PATCH 01/15] Write an env file directly during cmake build. Update
 scripts and readme docs to grab PYTHONPATH from env file.

Signed-off-by: zjgarvey <zjgarvey@gmail.com>
---
 CMakeLists.txt                                |  2 ++
 build_tools/ci/build_posix.sh                 |  1 +
 .../python_deploy/build_linux_packages.sh     |  3 +++
 build_tools/update_abstract_interp_lib.sh     | 26 +++++++------------
 build_tools/update_torch_ods.sh               | 25 +++++++-----------
 build_tools/write_env_file.sh                 | 26 -------------------
 docs/development.md                           | 19 ++++----------
 7 files changed, 31 insertions(+), 71 deletions(-)
 delete mode 100755 build_tools/write_env_file.sh

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 5413da9d9bab..a8275e9fd666 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -206,6 +206,8 @@ if(MLIR_ENABLE_BINDINGS_PYTHON)
   if(NOT TORCH_MLIR_PYTHON_PACKAGES_DIR)
     set(TORCH_MLIR_PYTHON_PACKAGES_DIR "${CMAKE_CURRENT_BINARY_DIR}/python_packages")
   endif()
+  file(WRITE "${TORCH_MLIR_SOURCE_DIR}/.env" "PYTHONPATH=${TORCH_MLIR_PYTHON_PACKAGES_DIR}/torch_mlir")
+  message(STATUS "PYTHONPATH to torch-mlir python_packages dir written to '${TORCH_MLIR_SOURCE_DIR}/.env'.")
 endif()
 
 add_subdirectory(test)
diff --git a/build_tools/ci/build_posix.sh b/build_tools/ci/build_posix.sh
index b9bb122acd37..69970927274c 100755
--- a/build_tools/ci/build_posix.sh
+++ b/build_tools/ci/build_posix.sh
@@ -41,6 +41,7 @@ cmake -S "$repo_root/externals/llvm-project/llvm" -B "$build_dir" \
   -GNinja \
   -DCMAKE_BUILD_TYPE=Release \
   -DPython3_EXECUTABLE="$(which python3)" \
+  -DPython_EXECUTABLE="$(which python3)" \
   -DLLVM_ENABLE_ASSERTIONS=ON \
   -DTORCH_MLIR_ENABLE_WERROR_FLAG=ON \
   -DCMAKE_INSTALL_PREFIX="$install_dir" \
diff --git a/build_tools/python_deploy/build_linux_packages.sh b/build_tools/python_deploy/build_linux_packages.sh
index b53f7482b839..e631a9ad4f5b 100755
--- a/build_tools/python_deploy/build_linux_packages.sh
+++ b/build_tools/python_deploy/build_linux_packages.sh
@@ -244,6 +244,7 @@ function build_in_tree() {
       -DTORCH_MLIR_SRC_PYTORCH_BRANCH=${TORCH_MLIR_SRC_PYTORCH_BRANCH} \
       -DTM_PYTORCH_INSTALL_WITHOUT_REBUILD=${TM_PYTORCH_INSTALL_WITHOUT_REBUILD} \
       -DPython3_EXECUTABLE="$(which python3)" \
+      -DPython_EXECUTABLE="$(which python3)" \
       /main_checkout/torch-mlir/externals/llvm-project/llvm
   cmake --build /main_checkout/torch-mlir/build --target tools/torch-mlir/all
   ccache -s
@@ -387,6 +388,7 @@ function build_out_of_tree() {
         -DLLVM_TARGETS_TO_BUILD=host \
         -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
         -DPython3_EXECUTABLE="$(which python3)" \
+        -DPython_EXECUTABLE="$(which python3)" \
         /main_checkout/torch-mlir/externals/llvm-project/llvm
     cmake --build /main_checkout/torch-mlir/llvm-build
   fi
@@ -409,6 +411,7 @@ function build_out_of_tree() {
       -DTORCH_MLIR_SRC_PYTORCH_BRANCH=${TORCH_MLIR_SRC_PYTORCH_BRANCH} \
       -DTM_PYTORCH_INSTALL_WITHOUT_REBUILD=${TM_PYTORCH_INSTALL_WITHOUT_REBUILD} \
       -DPython3_EXECUTABLE="$(which python3)" \
+      -DPython_EXECUTABLE="$(which python3)" \
       /main_checkout/torch-mlir
   cmake --build /main_checkout/torch-mlir/build_oot
   ccache -s
diff --git a/build_tools/update_abstract_interp_lib.sh b/build_tools/update_abstract_interp_lib.sh
index 4da20c3e715a..f5d606bdc81e 100755
--- a/build_tools/update_abstract_interp_lib.sh
+++ b/build_tools/update_abstract_interp_lib.sh
@@ -13,38 +13,32 @@
 set -euo pipefail
 
 src_dir="$(realpath "$(dirname "$0")"/..)"
-build_dir="$(realpath "${TORCH_MLIR_BUILD_DIR:-$src_dir/build}")"
 torch_transforms_cpp_dir="${src_dir}/lib/Dialect/Torch/Transforms"
 
-in_tree_pkg_dir="${build_dir}/tools/torch-mlir/python_packages"
-out_of_tree_pkg_dir="${build_dir}/python_packages"
+env_file_path="${src_dir}/.env"
 
-if [[ ! -d "${in_tree_pkg_dir}" && ! -d "${out_of_tree_pkg_dir}" ]]; then
-  echo "Couldn't find in-tree or out-of-tree build, exiting."
+if [[ ! -f "${env_file_path}" ]]; then
+  echo "Couldn't find an env file at ${env_file_path}!"
   exit 1
 fi
 
-# The `-nt` check works even if one of the two directories is missing.
-if [[ "${in_tree_pkg_dir}" -nt "${out_of_tree_pkg_dir}" ]]; then
-  python_packages_dir="${in_tree_pkg_dir}"
-else
-  python_packages_dir="${out_of_tree_pkg_dir}"
-fi
-
+# Get PYTHONPATH from env file.
+source $env_file_path
+# Update PYTHONPATH with externals if specified.
 TORCH_MLIR_EXT_PYTHONPATH="${TORCH_MLIR_EXT_PYTHONPATH:-""}"
-pypath="${python_packages_dir}/torch_mlir"
 if [ ! -z ${TORCH_MLIR_EXT_PYTHONPATH} ]; then
-  pypath="${pypath}:${TORCH_MLIR_EXT_PYTHONPATH}"
+  PYTHONPATH="${PYTHONPATH}:${TORCH_MLIR_EXT_PYTHONPATH}"
 fi
 TORCH_MLIR_EXT_MODULES="${TORCH_MLIR_EXT_MODULES:-""}"
+ext_module="${ext_module:-""}"
 if [ ! -z ${TORCH_MLIR_EXT_MODULES} ]; then
-  ext_module="${TORCH_MLIR_EXT_MODULES} "
+  ext_module="${TORCH_MLIR_EXT_MODULES}"
 fi
 
 # To enable this python package, manually build torch_mlir with:
 #   -DTORCH_MLIR_ENABLE_JIT_IR_IMPORTER=ON
 # TODO: move this package out of JIT_IR_IMPORTER.
-PYTHONPATH="${pypath}" python3 \
+python3 \
   -m torch_mlir.jit_ir_importer.build_tools.abstract_interp_lib_gen \
   --pytorch_op_extensions=${ext_module:-""} \
   --torch_transforms_cpp_dir="${torch_transforms_cpp_dir}"
diff --git a/build_tools/update_torch_ods.sh b/build_tools/update_torch_ods.sh
index efe0055d7e06..ce1c55222ae7 100755
--- a/build_tools/update_torch_ods.sh
+++ b/build_tools/update_torch_ods.sh
@@ -12,28 +12,21 @@
 set -euo pipefail
 
 src_dir="$(realpath "$(dirname "$0")"/..)"
-build_dir="$(realpath "${TORCH_MLIR_BUILD_DIR:-$src_dir/build}")"
 torch_ir_include_dir="${src_dir}/include/torch-mlir/Dialect/Torch/IR"
 
-in_tree_pkg_dir="${build_dir}/tools/torch-mlir/python_packages"
-out_of_tree_pkg_dir="${build_dir}/python_packages"
+env_file_path="${src_dir}/.env"
 
-if [[ ! -d "${in_tree_pkg_dir}" && ! -d "${out_of_tree_pkg_dir}" ]]; then
-  echo "Couldn't find in-tree or out-of-tree build, exiting."
+if [[ ! -f "${env_file_path}" ]]; then
+  echo "Couldn't find an env file at ${env_file_path}!"
   exit 1
 fi
 
-# The `-nt` check works even if one of the two directories is missing.
-if [[ "${in_tree_pkg_dir}" -nt "${out_of_tree_pkg_dir}" ]]; then
-  python_packages_dir="${in_tree_pkg_dir}"
-else
-  python_packages_dir="${out_of_tree_pkg_dir}"
-fi
-
+# Get PYTHONPATH from env file.
+source $env_file_path
+# Update PYTHONPATH with externals if specified.
 TORCH_MLIR_EXT_PYTHONPATH="${TORCH_MLIR_EXT_PYTHONPATH:-""}"
-pypath="${python_packages_dir}/torch_mlir"
 if [ ! -z ${TORCH_MLIR_EXT_PYTHONPATH} ]; then
-  pypath="${pypath}:${TORCH_MLIR_EXT_PYTHONPATH}"
+  PYTHONPATH="${PYTHONPATH}:${TORCH_MLIR_EXT_PYTHONPATH}"
 fi
 TORCH_MLIR_EXT_MODULES="${TORCH_MLIR_EXT_MODULES:-""}"
 ext_module="${ext_module:-""}"
@@ -44,8 +37,10 @@ fi
 set +u
 # To enable this python package, manually build torch_mlir with:
 #   -DTORCH_MLIR_ENABLE_PYTORCH_EXTENSIONS=ON
+#   -DTORCH_MLIR_ENABLE_JIT_IR_IMPORTER=ON
 # TODO: move this package out of JIT_IR_IMPORTER.
-PYTHONPATH="${PYTHONPATH}:${pypath}" python3 \
+echo $PYTHONPATH
+python3 \
   -m torch_mlir.jit_ir_importer.build_tools.torch_ods_gen \
   --torch_ir_include_dir="${torch_ir_include_dir}" \
   --pytorch_op_extensions="${ext_module}" \
diff --git a/build_tools/write_env_file.sh b/build_tools/write_env_file.sh
deleted file mode 100755
index 8f3c9a59357f..000000000000
--- a/build_tools/write_env_file.sh
+++ /dev/null
@@ -1,26 +0,0 @@
-#!/bin/bash
-# generate the .env file with default options.
-#
-# For arbitrary build/install directories, set the env variables:
-# - TORCH_MLIR_BUILD_DIR
-
-set -eu -o pipefail
-
-portable_realpath() {
-  # Create the directory if needed so that the `cd` doesn't fail.
-  mkdir -p "$1" && cd "$1" && pwd
-}
-
-td="$(portable_realpath "$(dirname "$0")"/..)"
-build_dir="$(portable_realpath "${TORCH_MLIR_BUILD_DIR:-$td/build}")"
-python_packages_dir="$build_dir/python_packages"
-
-write_env_file() {
-  echo "Updating $build_dir/.env file"
-  echo "PYTHONPATH=\"$(portable_realpath "$python_packages_dir/torch_mlir")\"" > "$build_dir/.env"
-  if ! cp "$build_dir/.env" "$td/.env"; then
-    echo "WARNING: Failed to write $td/.env"
-  fi
-}
-
-write_env_file
diff --git a/docs/development.md b/docs/development.md
index 360dff8f9df8..372d4b9c4bf4 100644
--- a/docs/development.md
+++ b/docs/development.md
@@ -197,13 +197,14 @@ TIP: add multiple target options to stack build phases
 #### Linux and macOS
 
 ```shell
-export PYTHONPATH=`pwd`/build/python_packages/torch_mlir:`pwd`/test/python/fx_importer
+source $PWD/.env && export PYTHONPATH="${PYTHONPATH}:${PWD}/test/python/fx_importer"
 ```
 
 #### Windows PowerShell
 
 ```shell
-$env:PYTHONPATH = "$PWD/build/tools/torch-mlir/python_packages/torch_mlir;$PWD/test/python/fx_importer"
+$env:PYTHONPATH = ... # Get from .env file in torch-mlir repo root. 
+$env:PYTHONPATH += ";${PWD}/test/python/fx_importer"
 ```
 
 ### Testing MLIR output in various dialects
@@ -214,8 +215,6 @@ Make sure you have activated the virtualenv and set the `PYTHONPATH` above
 (if running on Windows, modify the environment variable as shown above):
 
 ```shell
-source mlir_venv/bin/activate
-export PYTHONPATH=`pwd`/build/tools/torch-mlir/python_packages/torch_mlir:`pwd`/test/python/fx_importer
 python test/python/fx_importer/basic_test.py
 ```
 
@@ -226,10 +225,10 @@ using torchscript with the example `projects/pt1/examples/torchscript_resnet18_a
 This path doesn't give access to the current generation work that is being driven via the fx_importer
 and may lead to errors.
 
-Same as above, but with different python path and example:
+Same as above, but you need to append to the PYTHONPATH:
 
 ```shell
-export PYTHONPATH=`pwd`/build/tools/torch-mlir/python_packages/torch_mlir:`pwd`/projects/pt1/examples
+export PYTHONPATH="${PYTHONPATH}:/projects/pt1/examples"
 python projects/pt1/examples/torchscript_resnet18_all_output_types.py
 ```
 
@@ -259,14 +258,6 @@ jupyter notebook
 [Example IR](https://gist.github.com/silvasean/e74780f8a8a449339aac05c51e8b0caa) for a simple 1 layer MLP to show the compilation steps from TorchScript.
 
 
-### Interactive Use
-
-The `build_tools/write_env_file.sh` script will output a `.env`
-file in the workspace folder with the correct PYTHONPATH set. This allows
-tools like VSCode to work by default for debugging. This file can also be
-manually `source`'d in a shell.
-
-
 ### Bazel Build
 
 > **NOTE** Our Bazel build follows LLVM's Bazel build policy: only the

From 493696f635a3b4ea9d5f77d46ee5d3f5183397cd Mon Sep 17 00:00:00 2001
From: zjgarvey <zjgarvey@gmail.com>
Date: Fri, 3 Oct 2025 13:34:40 -0700
Subject: [PATCH 02/15] Remove debug print

Signed-off-by: zjgarvey <zjgarvey@gmail.com>
---
 build_tools/update_torch_ods.sh | 1 -
 1 file changed, 1 deletion(-)

diff --git a/build_tools/update_torch_ods.sh b/build_tools/update_torch_ods.sh
index ce1c55222ae7..b09db9525aad 100755
--- a/build_tools/update_torch_ods.sh
+++ b/build_tools/update_torch_ods.sh
@@ -39,7 +39,6 @@ set +u
 #   -DTORCH_MLIR_ENABLE_PYTORCH_EXTENSIONS=ON
 #   -DTORCH_MLIR_ENABLE_JIT_IR_IMPORTER=ON
 # TODO: move this package out of JIT_IR_IMPORTER.
-echo $PYTHONPATH
 python3 \
   -m torch_mlir.jit_ir_importer.build_tools.torch_ods_gen \
   --torch_ir_include_dir="${torch_ir_include_dir}" \

From 69783520276ed4a9a14dc9aea5888c062390e218 Mon Sep 17 00:00:00 2001
From: zjgarvey <zjgarvey@gmail.com>
Date: Fri, 3 Oct 2025 13:37:13 -0700
Subject: [PATCH 03/15] remove additional mention of env file writing script

Signed-off-by: zjgarvey <zjgarvey@gmail.com>
---
 docs/development.md | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/docs/development.md b/docs/development.md
index 372d4b9c4bf4..f41dc5a59274 100644
--- a/docs/development.md
+++ b/docs/development.md
@@ -469,10 +469,6 @@ Torch-MLIR has two types of tests:
 
 ### Running execution (end-to-end) tests:
 
-> **Note**
-> An `.env` file must be generated via `build_tools/write_env_file.sh` before these commands can be run.
-
-
 The following assumes you are in the `projects/pt1`  directory:
 
 ```shell

From 3c9a1f84de0346d8dc36b627e633712f7ff3443d Mon Sep 17 00:00:00 2001
From: zjgarvey <zjgarvey@gmail.com>
Date: Fri, 3 Oct 2025 13:39:15 -0700
Subject: [PATCH 04/15] lint

Signed-off-by: zjgarvey <zjgarvey@gmail.com>
---
 docs/development.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/development.md b/docs/development.md
index f41dc5a59274..d1d330720150 100644
--- a/docs/development.md
+++ b/docs/development.md
@@ -203,7 +203,7 @@ source $PWD/.env && export PYTHONPATH="${PYTHONPATH}:${PWD}/test/python/fx_impor
 #### Windows PowerShell
 
 ```shell
-$env:PYTHONPATH = ... # Get from .env file in torch-mlir repo root. 
+$env:PYTHONPATH = ... # Get from .env file in torch-mlir repo root.
 $env:PYTHONPATH += ";${PWD}/test/python/fx_importer"
 ```
 

From 09de7e5686bf288206ac8a8efdb844a056070256 Mon Sep 17 00:00:00 2001
From: zjgarvey <zjgarvey@gmail.com>
Date: Fri, 3 Oct 2025 14:01:04 -0700
Subject: [PATCH 05/15] Print PYTHONPATH for debugging

Signed-off-by: zjgarvey <zjgarvey@gmail.com>
---
 build_tools/update_abstract_interp_lib.sh | 4 +++-
 build_tools/update_torch_ods.sh           | 2 ++
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/build_tools/update_abstract_interp_lib.sh b/build_tools/update_abstract_interp_lib.sh
index f5d606bdc81e..f44f74b906f0 100755
--- a/build_tools/update_abstract_interp_lib.sh
+++ b/build_tools/update_abstract_interp_lib.sh
@@ -35,10 +35,12 @@ if [ ! -z ${TORCH_MLIR_EXT_MODULES} ]; then
   ext_module="${TORCH_MLIR_EXT_MODULES}"
 fi
 
+echo "Running python tool with PYTHONPATH=${PYTHONPATH}"
+
 # To enable this python package, manually build torch_mlir with:
 #   -DTORCH_MLIR_ENABLE_JIT_IR_IMPORTER=ON
 # TODO: move this package out of JIT_IR_IMPORTER.
 python3 \
   -m torch_mlir.jit_ir_importer.build_tools.abstract_interp_lib_gen \
-  --pytorch_op_extensions=${ext_module:-""} \
+  --pytorch_op_extensions=${ext_module} \
   --torch_transforms_cpp_dir="${torch_transforms_cpp_dir}"
diff --git a/build_tools/update_torch_ods.sh b/build_tools/update_torch_ods.sh
index b09db9525aad..300019b5ec76 100755
--- a/build_tools/update_torch_ods.sh
+++ b/build_tools/update_torch_ods.sh
@@ -34,6 +34,8 @@ if [ ! -z ${TORCH_MLIR_EXT_MODULES} ]; then
   ext_module="${TORCH_MLIR_EXT_MODULES}"
 fi
 
+echo "Running python tool with PYTHONPATH=${PYTHONPATH}"
+
 set +u
 # To enable this python package, manually build torch_mlir with:
 #   -DTORCH_MLIR_ENABLE_PYTORCH_EXTENSIONS=ON

From 54a63f8a4fc4fe51a75a51beb4141d3453eed3b7 Mon Sep 17 00:00:00 2001
From: zjgarvey <zjgarvey@gmail.com>
Date: Fri, 3 Oct 2025 14:47:41 -0700
Subject: [PATCH 06/15] Set PYTHONPATH in tool run commands and remove debug
 statement.

Signed-off-by: zjgarvey <zjgarvey@gmail.com>
---
 build_tools/update_abstract_interp_lib.sh | 4 +---
 build_tools/update_torch_ods.sh           | 5 +----
 2 files changed, 2 insertions(+), 7 deletions(-)

diff --git a/build_tools/update_abstract_interp_lib.sh b/build_tools/update_abstract_interp_lib.sh
index f44f74b906f0..9b973e1e907a 100755
--- a/build_tools/update_abstract_interp_lib.sh
+++ b/build_tools/update_abstract_interp_lib.sh
@@ -35,12 +35,10 @@ if [ ! -z ${TORCH_MLIR_EXT_MODULES} ]; then
   ext_module="${TORCH_MLIR_EXT_MODULES}"
 fi
 
-echo "Running python tool with PYTHONPATH=${PYTHONPATH}"
-
 # To enable this python package, manually build torch_mlir with:
 #   -DTORCH_MLIR_ENABLE_JIT_IR_IMPORTER=ON
 # TODO: move this package out of JIT_IR_IMPORTER.
-python3 \
+PYTHONPATH=${PYTHONPATH} python3 \
   -m torch_mlir.jit_ir_importer.build_tools.abstract_interp_lib_gen \
   --pytorch_op_extensions=${ext_module} \
   --torch_transforms_cpp_dir="${torch_transforms_cpp_dir}"
diff --git a/build_tools/update_torch_ods.sh b/build_tools/update_torch_ods.sh
index 300019b5ec76..3760ac96ea12 100755
--- a/build_tools/update_torch_ods.sh
+++ b/build_tools/update_torch_ods.sh
@@ -34,14 +34,11 @@ if [ ! -z ${TORCH_MLIR_EXT_MODULES} ]; then
   ext_module="${TORCH_MLIR_EXT_MODULES}"
 fi
 
-echo "Running python tool with PYTHONPATH=${PYTHONPATH}"
-
-set +u
 # To enable this python package, manually build torch_mlir with:
 #   -DTORCH_MLIR_ENABLE_PYTORCH_EXTENSIONS=ON
 #   -DTORCH_MLIR_ENABLE_JIT_IR_IMPORTER=ON
 # TODO: move this package out of JIT_IR_IMPORTER.
-python3 \
+PYTHONPATH=${PYTHONPATH} python3 \
   -m torch_mlir.jit_ir_importer.build_tools.torch_ods_gen \
   --torch_ir_include_dir="${torch_ir_include_dir}" \
   --pytorch_op_extensions="${ext_module}" \

From 884a062d8175bff08b8afc0517c2b2fd4dad43a1 Mon Sep 17 00:00:00 2001
From: zjgarvey <zjgarvey@gmail.com>
Date: Mon, 6 Oct 2025 09:09:39 -0700
Subject: [PATCH 07/15] Add DPython configure flag for cmake commands.

Signed-off-by: zjgarvey <zjgarvey@gmail.com>
---
 build_tools/ci/build_posix.sh                     | 1 +
 build_tools/python_deploy/build_linux_packages.sh | 3 +++
 2 files changed, 4 insertions(+)

diff --git a/build_tools/ci/build_posix.sh b/build_tools/ci/build_posix.sh
index b9bb122acd37..69970927274c 100755
--- a/build_tools/ci/build_posix.sh
+++ b/build_tools/ci/build_posix.sh
@@ -41,6 +41,7 @@ cmake -S "$repo_root/externals/llvm-project/llvm" -B "$build_dir" \
   -GNinja \
   -DCMAKE_BUILD_TYPE=Release \
   -DPython3_EXECUTABLE="$(which python3)" \
+  -DPython_EXECUTABLE="$(which python3)" \
   -DLLVM_ENABLE_ASSERTIONS=ON \
   -DTORCH_MLIR_ENABLE_WERROR_FLAG=ON \
   -DCMAKE_INSTALL_PREFIX="$install_dir" \
diff --git a/build_tools/python_deploy/build_linux_packages.sh b/build_tools/python_deploy/build_linux_packages.sh
index b53f7482b839..e631a9ad4f5b 100755
--- a/build_tools/python_deploy/build_linux_packages.sh
+++ b/build_tools/python_deploy/build_linux_packages.sh
@@ -244,6 +244,7 @@ function build_in_tree() {
       -DTORCH_MLIR_SRC_PYTORCH_BRANCH=${TORCH_MLIR_SRC_PYTORCH_BRANCH} \
       -DTM_PYTORCH_INSTALL_WITHOUT_REBUILD=${TM_PYTORCH_INSTALL_WITHOUT_REBUILD} \
       -DPython3_EXECUTABLE="$(which python3)" \
+      -DPython_EXECUTABLE="$(which python3)" \
       /main_checkout/torch-mlir/externals/llvm-project/llvm
   cmake --build /main_checkout/torch-mlir/build --target tools/torch-mlir/all
   ccache -s
@@ -387,6 +388,7 @@ function build_out_of_tree() {
         -DLLVM_TARGETS_TO_BUILD=host \
         -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
         -DPython3_EXECUTABLE="$(which python3)" \
+        -DPython_EXECUTABLE="$(which python3)" \
         /main_checkout/torch-mlir/externals/llvm-project/llvm
     cmake --build /main_checkout/torch-mlir/llvm-build
   fi
@@ -409,6 +411,7 @@ function build_out_of_tree() {
       -DTORCH_MLIR_SRC_PYTORCH_BRANCH=${TORCH_MLIR_SRC_PYTORCH_BRANCH} \
       -DTM_PYTORCH_INSTALL_WITHOUT_REBUILD=${TM_PYTORCH_INSTALL_WITHOUT_REBUILD} \
       -DPython3_EXECUTABLE="$(which python3)" \
+      -DPython_EXECUTABLE="$(which python3)" \
       /main_checkout/torch-mlir
   cmake --build /main_checkout/torch-mlir/build_oot
   ccache -s

From 3e943b6a9249e6738ff4bb2d6ffbacb9c70b81a3 Mon Sep 17 00:00:00 2001
From: zjgarvey <zjgarvey@gmail.com>
Date: Mon, 6 Oct 2025 09:32:54 -0700
Subject: [PATCH 08/15] Update `write_env_file.sh` to correctly find either
 in-tree or out-of-tree dir.

Signed-off-by: zjgarvey <zjgarvey@gmail.com>
---
 build_tools/write_env_file.sh | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/build_tools/write_env_file.sh b/build_tools/write_env_file.sh
index 8f3c9a59357f..83a0bcac4ef9 100755
--- a/build_tools/write_env_file.sh
+++ b/build_tools/write_env_file.sh
@@ -13,7 +13,21 @@ portable_realpath() {
 
 td="$(portable_realpath "$(dirname "$0")"/..)"
 build_dir="$(portable_realpath "${TORCH_MLIR_BUILD_DIR:-$td/build}")"
-python_packages_dir="$build_dir/python_packages"
+
+in_tree_pkg_dir="${build_dir}/tools/torch-mlir/python_packages"
+out_of_tree_pkg_dir="${build_dir}/python_packages"
+
+if [[ ! -d "${in_tree_pkg_dir}" && ! -d "${out_of_tree_pkg_dir}" ]]; then
+  echo "Couldn't find in-tree or out-of-tree build, exiting."
+  exit 1
+fi
+
+# The `-nt` check works even if one of the two directories is missing.
+if [[ "${in_tree_pkg_dir}" -nt "${out_of_tree_pkg_dir}" ]]; then
+  python_packages_dir="${in_tree_pkg_dir}"
+else
+  python_packages_dir="${out_of_tree_pkg_dir}"
+fi
 
 write_env_file() {
   echo "Updating $build_dir/.env file"

From 41fdb381d9a08e2b2c1067bedc34d09835eb2c13 Mon Sep 17 00:00:00 2001
From: zjgarvey <zjgarvey@gmail.com>
Date: Mon, 6 Oct 2025 09:33:21 -0700
Subject: [PATCH 09/15] Update development.md with better instructions for
 PYTHONPATH handling.

Signed-off-by: zjgarvey <zjgarvey@gmail.com>
---
 docs/development.md | 40 +++++++++++++++++++++++++++++-----------
 1 file changed, 29 insertions(+), 11 deletions(-)

diff --git a/docs/development.md b/docs/development.md
index 360dff8f9df8..1c85437376b1 100644
--- a/docs/development.md
+++ b/docs/development.md
@@ -194,16 +194,40 @@ TIP: add multiple target options to stack build phases
 
 ### Setup Python Environment to export the built Python packages
 
+When CMake is configured with `-DMLIR_ENABLE_BINDINGS_PYTHON=ON`, the python packages will typically be located in either:
+
+1. `./build/tools/torch-mlir/python_packages/`  if doing an in-tree build.
+2. `./build/python_packages/` if doing an out-of-tree build.
+
+For the following sections, let `python_pkg_dir` represent whichever of the above is relevant for your build setup. On Linux and macOS, you can run `./build_tools/write_env_file.sh` to generate a file `./.env` in your root source directory with the correct `PYTHONPATH`.
+
 #### Linux and macOS
 
+To get the base `PYTHONPATH`, run:
+
 ```shell
-export PYTHONPATH=`pwd`/build/python_packages/torch_mlir:`pwd`/test/python/fx_importer
+./build_tools/write_env_file.sh
+source ./.env && export PYTHONPATH
+```
+
+To run fx_importer tests, you can append the following:
+
+```
+export PYTHONPATH="${PYTHONPATH}":/test/python/fx_importer"
 ```
 
 #### Windows PowerShell
 
+To get the base `PYTHONPATH`, identify your `python_pkg_dir` from above and set this variable in your environment:
+
+```shell
+$env:PYTHONPATH = "<python_pkg_dir>/torch-mlir"
+```
+
+To run fx_importer tests, you can append the following:
+
 ```shell
-$env:PYTHONPATH = "$PWD/build/tools/torch-mlir/python_packages/torch_mlir;$PWD/test/python/fx_importer"
+$env:PYTHONPATH += ";$PWD/test/python/fx_importer"
 ```
 
 ### Testing MLIR output in various dialects
@@ -211,13 +235,7 @@ $env:PYTHONPATH = "$PWD/build/tools/torch-mlir/python_packages/torch_mlir;$PWD/t
 To test the MLIR output to torch dialect, you can use `test/python/fx_importer/basic_test.py`.
 
 Make sure you have activated the virtualenv and set the `PYTHONPATH` above
-(if running on Windows, modify the environment variable as shown above):
-
-```shell
-source mlir_venv/bin/activate
-export PYTHONPATH=`pwd`/build/tools/torch-mlir/python_packages/torch_mlir:`pwd`/test/python/fx_importer
-python test/python/fx_importer/basic_test.py
-```
+(if running on Windows, modify the environment variable as shown above).
 
 This will display the basic example in TORCH dialect.
 
@@ -226,10 +244,10 @@ using torchscript with the example `projects/pt1/examples/torchscript_resnet18_a
 This path doesn't give access to the current generation work that is being driven via the fx_importer
 and may lead to errors.
 
-Same as above, but with different python path and example:
+The base `PYTHONPATH` should be set as above, then the example can be run with the following command (similar on Windows):
 
 ```shell
-export PYTHONPATH=`pwd`/build/tools/torch-mlir/python_packages/torch_mlir:`pwd`/projects/pt1/examples
+export PYTHONPATH="${PYTHONPATH}:$PWD/projects/pt1/examples"
 python projects/pt1/examples/torchscript_resnet18_all_output_types.py
 ```
 

From 55446dccedd3eb3b82cbed9a089615faaace9e1f Mon Sep 17 00:00:00 2001
From: zjgarvey <zjgarvey@gmail.com>
Date: Mon, 6 Oct 2025 11:46:08 -0700
Subject: [PATCH 10/15] Revert "Set PYTHONPATH in tool run commands and remove
 debug statement."

This reverts commit 54a63f8a4fc4fe51a75a51beb4141d3453eed3b7.
---
 build_tools/update_abstract_interp_lib.sh | 4 +++-
 build_tools/update_torch_ods.sh           | 5 ++++-
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/build_tools/update_abstract_interp_lib.sh b/build_tools/update_abstract_interp_lib.sh
index 9b973e1e907a..f44f74b906f0 100755
--- a/build_tools/update_abstract_interp_lib.sh
+++ b/build_tools/update_abstract_interp_lib.sh
@@ -35,10 +35,12 @@ if [ ! -z ${TORCH_MLIR_EXT_MODULES} ]; then
   ext_module="${TORCH_MLIR_EXT_MODULES}"
 fi
 
+echo "Running python tool with PYTHONPATH=${PYTHONPATH}"
+
 # To enable this python package, manually build torch_mlir with:
 #   -DTORCH_MLIR_ENABLE_JIT_IR_IMPORTER=ON
 # TODO: move this package out of JIT_IR_IMPORTER.
-PYTHONPATH=${PYTHONPATH} python3 \
+python3 \
   -m torch_mlir.jit_ir_importer.build_tools.abstract_interp_lib_gen \
   --pytorch_op_extensions=${ext_module} \
   --torch_transforms_cpp_dir="${torch_transforms_cpp_dir}"
diff --git a/build_tools/update_torch_ods.sh b/build_tools/update_torch_ods.sh
index 3760ac96ea12..300019b5ec76 100755
--- a/build_tools/update_torch_ods.sh
+++ b/build_tools/update_torch_ods.sh
@@ -34,11 +34,14 @@ if [ ! -z ${TORCH_MLIR_EXT_MODULES} ]; then
   ext_module="${TORCH_MLIR_EXT_MODULES}"
 fi
 
+echo "Running python tool with PYTHONPATH=${PYTHONPATH}"
+
+set +u
 # To enable this python package, manually build torch_mlir with:
 #   -DTORCH_MLIR_ENABLE_PYTORCH_EXTENSIONS=ON
 #   -DTORCH_MLIR_ENABLE_JIT_IR_IMPORTER=ON
 # TODO: move this package out of JIT_IR_IMPORTER.
-PYTHONPATH=${PYTHONPATH} python3 \
+python3 \
   -m torch_mlir.jit_ir_importer.build_tools.torch_ods_gen \
   --torch_ir_include_dir="${torch_ir_include_dir}" \
   --pytorch_op_extensions="${ext_module}" \

From 063d26a89ab29acbd94cd5d6009548a245203069 Mon Sep 17 00:00:00 2001
From: zjgarvey <zjgarvey@gmail.com>
Date: Mon, 6 Oct 2025 11:46:17 -0700
Subject: [PATCH 11/15] Revert "Print PYTHONPATH for debugging"

This reverts commit 09de7e5686bf288206ac8a8efdb844a056070256.
---
 build_tools/update_abstract_interp_lib.sh | 4 +---
 build_tools/update_torch_ods.sh           | 2 --
 2 files changed, 1 insertion(+), 5 deletions(-)

diff --git a/build_tools/update_abstract_interp_lib.sh b/build_tools/update_abstract_interp_lib.sh
index f44f74b906f0..f5d606bdc81e 100755
--- a/build_tools/update_abstract_interp_lib.sh
+++ b/build_tools/update_abstract_interp_lib.sh
@@ -35,12 +35,10 @@ if [ ! -z ${TORCH_MLIR_EXT_MODULES} ]; then
   ext_module="${TORCH_MLIR_EXT_MODULES}"
 fi
 
-echo "Running python tool with PYTHONPATH=${PYTHONPATH}"
-
 # To enable this python package, manually build torch_mlir with:
 #   -DTORCH_MLIR_ENABLE_JIT_IR_IMPORTER=ON
 # TODO: move this package out of JIT_IR_IMPORTER.
 python3 \
   -m torch_mlir.jit_ir_importer.build_tools.abstract_interp_lib_gen \
-  --pytorch_op_extensions=${ext_module} \
+  --pytorch_op_extensions=${ext_module:-""} \
   --torch_transforms_cpp_dir="${torch_transforms_cpp_dir}"
diff --git a/build_tools/update_torch_ods.sh b/build_tools/update_torch_ods.sh
index 300019b5ec76..b09db9525aad 100755
--- a/build_tools/update_torch_ods.sh
+++ b/build_tools/update_torch_ods.sh
@@ -34,8 +34,6 @@ if [ ! -z ${TORCH_MLIR_EXT_MODULES} ]; then
   ext_module="${TORCH_MLIR_EXT_MODULES}"
 fi
 
-echo "Running python tool with PYTHONPATH=${PYTHONPATH}"
-
 set +u
 # To enable this python package, manually build torch_mlir with:
 #   -DTORCH_MLIR_ENABLE_PYTORCH_EXTENSIONS=ON

From c14a5a20d50b1affa261d648cc08150a98f4c3ac Mon Sep 17 00:00:00 2001
From: zjgarvey <zjgarvey@gmail.com>
Date: Mon, 6 Oct 2025 11:46:19 -0700
Subject: [PATCH 12/15] Revert "lint"

This reverts commit 3c9a1f84de0346d8dc36b627e633712f7ff3443d.
---
 docs/development.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/development.md b/docs/development.md
index d1d330720150..f41dc5a59274 100644
--- a/docs/development.md
+++ b/docs/development.md
@@ -203,7 +203,7 @@ source $PWD/.env && export PYTHONPATH="${PYTHONPATH}:${PWD}/test/python/fx_impor
 #### Windows PowerShell
 
 ```shell
-$env:PYTHONPATH = ... # Get from .env file in torch-mlir repo root.
+$env:PYTHONPATH = ... # Get from .env file in torch-mlir repo root. 
 $env:PYTHONPATH += ";${PWD}/test/python/fx_importer"
 ```
 

From a86a4eb1cc820d1ef02dd974bfbc59245a15daeb Mon Sep 17 00:00:00 2001
From: zjgarvey <zjgarvey@gmail.com>
Date: Mon, 6 Oct 2025 11:46:22 -0700
Subject: [PATCH 13/15] Revert "remove additional mention of env file writing
 script"

This reverts commit 69783520276ed4a9a14dc9aea5888c062390e218.
---
 docs/development.md | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/docs/development.md b/docs/development.md
index f41dc5a59274..372d4b9c4bf4 100644
--- a/docs/development.md
+++ b/docs/development.md
@@ -469,6 +469,10 @@ Torch-MLIR has two types of tests:
 
 ### Running execution (end-to-end) tests:
 
+> **Note**
+> An `.env` file must be generated via `build_tools/write_env_file.sh` before these commands can be run.
+
+
 The following assumes you are in the `projects/pt1`  directory:
 
 ```shell

From ad97a18954129605c04169fb7fb539ee0ff7fd71 Mon Sep 17 00:00:00 2001
From: zjgarvey <zjgarvey@gmail.com>
Date: Mon, 6 Oct 2025 11:46:24 -0700
Subject: [PATCH 14/15] Revert "Remove debug print"

This reverts commit 493696f635a3b4ea9d5f77d46ee5d3f5183397cd.
---
 build_tools/update_torch_ods.sh | 1 +
 1 file changed, 1 insertion(+)

diff --git a/build_tools/update_torch_ods.sh b/build_tools/update_torch_ods.sh
index b09db9525aad..ce1c55222ae7 100755
--- a/build_tools/update_torch_ods.sh
+++ b/build_tools/update_torch_ods.sh
@@ -39,6 +39,7 @@ set +u
 #   -DTORCH_MLIR_ENABLE_PYTORCH_EXTENSIONS=ON
 #   -DTORCH_MLIR_ENABLE_JIT_IR_IMPORTER=ON
 # TODO: move this package out of JIT_IR_IMPORTER.
+echo $PYTHONPATH
 python3 \
   -m torch_mlir.jit_ir_importer.build_tools.torch_ods_gen \
   --torch_ir_include_dir="${torch_ir_include_dir}" \

From cf8a615ae97826770d452cfad13c6646b0741a5d Mon Sep 17 00:00:00 2001
From: zjgarvey <zjgarvey@gmail.com>
Date: Mon, 6 Oct 2025 11:46:25 -0700
Subject: [PATCH 15/15] Revert "Write an env file directly during cmake build."

This reverts commit 8fde78e4c2bacf07b76a9aadb57e03b807c5e85d.
---
 CMakeLists.txt                                |  2 --
 build_tools/ci/build_posix.sh                 |  1 -
 .../python_deploy/build_linux_packages.sh     |  3 ---
 build_tools/update_abstract_interp_lib.sh     | 26 ++++++++++++-------
 build_tools/update_torch_ods.sh               | 25 +++++++++++-------
 build_tools/write_env_file.sh                 | 26 +++++++++++++++++++
 docs/development.md                           | 19 ++++++++++----
 7 files changed, 71 insertions(+), 31 deletions(-)
 create mode 100755 build_tools/write_env_file.sh

diff --git a/CMakeLists.txt b/CMakeLists.txt
index a8275e9fd666..5413da9d9bab 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -206,8 +206,6 @@ if(MLIR_ENABLE_BINDINGS_PYTHON)
   if(NOT TORCH_MLIR_PYTHON_PACKAGES_DIR)
     set(TORCH_MLIR_PYTHON_PACKAGES_DIR "${CMAKE_CURRENT_BINARY_DIR}/python_packages")
   endif()
-  file(WRITE "${TORCH_MLIR_SOURCE_DIR}/.env" "PYTHONPATH=${TORCH_MLIR_PYTHON_PACKAGES_DIR}/torch_mlir")
-  message(STATUS "PYTHONPATH to torch-mlir python_packages dir written to '${TORCH_MLIR_SOURCE_DIR}/.env'.")
 endif()
 
 add_subdirectory(test)
diff --git a/build_tools/ci/build_posix.sh b/build_tools/ci/build_posix.sh
index 69970927274c..b9bb122acd37 100755
--- a/build_tools/ci/build_posix.sh
+++ b/build_tools/ci/build_posix.sh
@@ -41,7 +41,6 @@ cmake -S "$repo_root/externals/llvm-project/llvm" -B "$build_dir" \
   -GNinja \
   -DCMAKE_BUILD_TYPE=Release \
   -DPython3_EXECUTABLE="$(which python3)" \
-  -DPython_EXECUTABLE="$(which python3)" \
   -DLLVM_ENABLE_ASSERTIONS=ON \
   -DTORCH_MLIR_ENABLE_WERROR_FLAG=ON \
   -DCMAKE_INSTALL_PREFIX="$install_dir" \
diff --git a/build_tools/python_deploy/build_linux_packages.sh b/build_tools/python_deploy/build_linux_packages.sh
index e631a9ad4f5b..b53f7482b839 100755
--- a/build_tools/python_deploy/build_linux_packages.sh
+++ b/build_tools/python_deploy/build_linux_packages.sh
@@ -244,7 +244,6 @@ function build_in_tree() {
       -DTORCH_MLIR_SRC_PYTORCH_BRANCH=${TORCH_MLIR_SRC_PYTORCH_BRANCH} \
       -DTM_PYTORCH_INSTALL_WITHOUT_REBUILD=${TM_PYTORCH_INSTALL_WITHOUT_REBUILD} \
       -DPython3_EXECUTABLE="$(which python3)" \
-      -DPython_EXECUTABLE="$(which python3)" \
       /main_checkout/torch-mlir/externals/llvm-project/llvm
   cmake --build /main_checkout/torch-mlir/build --target tools/torch-mlir/all
   ccache -s
@@ -388,7 +387,6 @@ function build_out_of_tree() {
         -DLLVM_TARGETS_TO_BUILD=host \
         -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
         -DPython3_EXECUTABLE="$(which python3)" \
-        -DPython_EXECUTABLE="$(which python3)" \
         /main_checkout/torch-mlir/externals/llvm-project/llvm
     cmake --build /main_checkout/torch-mlir/llvm-build
   fi
@@ -411,7 +409,6 @@ function build_out_of_tree() {
       -DTORCH_MLIR_SRC_PYTORCH_BRANCH=${TORCH_MLIR_SRC_PYTORCH_BRANCH} \
       -DTM_PYTORCH_INSTALL_WITHOUT_REBUILD=${TM_PYTORCH_INSTALL_WITHOUT_REBUILD} \
       -DPython3_EXECUTABLE="$(which python3)" \
-      -DPython_EXECUTABLE="$(which python3)" \
       /main_checkout/torch-mlir
   cmake --build /main_checkout/torch-mlir/build_oot
   ccache -s
diff --git a/build_tools/update_abstract_interp_lib.sh b/build_tools/update_abstract_interp_lib.sh
index f5d606bdc81e..4da20c3e715a 100755
--- a/build_tools/update_abstract_interp_lib.sh
+++ b/build_tools/update_abstract_interp_lib.sh
@@ -13,32 +13,38 @@
 set -euo pipefail
 
 src_dir="$(realpath "$(dirname "$0")"/..)"
+build_dir="$(realpath "${TORCH_MLIR_BUILD_DIR:-$src_dir/build}")"
 torch_transforms_cpp_dir="${src_dir}/lib/Dialect/Torch/Transforms"
 
-env_file_path="${src_dir}/.env"
+in_tree_pkg_dir="${build_dir}/tools/torch-mlir/python_packages"
+out_of_tree_pkg_dir="${build_dir}/python_packages"
 
-if [[ ! -f "${env_file_path}" ]]; then
-  echo "Couldn't find an env file at ${env_file_path}!"
+if [[ ! -d "${in_tree_pkg_dir}" && ! -d "${out_of_tree_pkg_dir}" ]]; then
+  echo "Couldn't find in-tree or out-of-tree build, exiting."
   exit 1
 fi
 
-# Get PYTHONPATH from env file.
-source $env_file_path
-# Update PYTHONPATH with externals if specified.
+# The `-nt` check works even if one of the two directories is missing.
+if [[ "${in_tree_pkg_dir}" -nt "${out_of_tree_pkg_dir}" ]]; then
+  python_packages_dir="${in_tree_pkg_dir}"
+else
+  python_packages_dir="${out_of_tree_pkg_dir}"
+fi
+
 TORCH_MLIR_EXT_PYTHONPATH="${TORCH_MLIR_EXT_PYTHONPATH:-""}"
+pypath="${python_packages_dir}/torch_mlir"
 if [ ! -z ${TORCH_MLIR_EXT_PYTHONPATH} ]; then
-  PYTHONPATH="${PYTHONPATH}:${TORCH_MLIR_EXT_PYTHONPATH}"
+  pypath="${pypath}:${TORCH_MLIR_EXT_PYTHONPATH}"
 fi
 TORCH_MLIR_EXT_MODULES="${TORCH_MLIR_EXT_MODULES:-""}"
-ext_module="${ext_module:-""}"
 if [ ! -z ${TORCH_MLIR_EXT_MODULES} ]; then
-  ext_module="${TORCH_MLIR_EXT_MODULES}"
+  ext_module="${TORCH_MLIR_EXT_MODULES} "
 fi
 
 # To enable this python package, manually build torch_mlir with:
 #   -DTORCH_MLIR_ENABLE_JIT_IR_IMPORTER=ON
 # TODO: move this package out of JIT_IR_IMPORTER.
-python3 \
+PYTHONPATH="${pypath}" python3 \
   -m torch_mlir.jit_ir_importer.build_tools.abstract_interp_lib_gen \
   --pytorch_op_extensions=${ext_module:-""} \
   --torch_transforms_cpp_dir="${torch_transforms_cpp_dir}"
diff --git a/build_tools/update_torch_ods.sh b/build_tools/update_torch_ods.sh
index ce1c55222ae7..efe0055d7e06 100755
--- a/build_tools/update_torch_ods.sh
+++ b/build_tools/update_torch_ods.sh
@@ -12,21 +12,28 @@
 set -euo pipefail
 
 src_dir="$(realpath "$(dirname "$0")"/..)"
+build_dir="$(realpath "${TORCH_MLIR_BUILD_DIR:-$src_dir/build}")"
 torch_ir_include_dir="${src_dir}/include/torch-mlir/Dialect/Torch/IR"
 
-env_file_path="${src_dir}/.env"
+in_tree_pkg_dir="${build_dir}/tools/torch-mlir/python_packages"
+out_of_tree_pkg_dir="${build_dir}/python_packages"
 
-if [[ ! -f "${env_file_path}" ]]; then
-  echo "Couldn't find an env file at ${env_file_path}!"
+if [[ ! -d "${in_tree_pkg_dir}" && ! -d "${out_of_tree_pkg_dir}" ]]; then
+  echo "Couldn't find in-tree or out-of-tree build, exiting."
   exit 1
 fi
 
-# Get PYTHONPATH from env file.
-source $env_file_path
-# Update PYTHONPATH with externals if specified.
+# The `-nt` check works even if one of the two directories is missing.
+if [[ "${in_tree_pkg_dir}" -nt "${out_of_tree_pkg_dir}" ]]; then
+  python_packages_dir="${in_tree_pkg_dir}"
+else
+  python_packages_dir="${out_of_tree_pkg_dir}"
+fi
+
 TORCH_MLIR_EXT_PYTHONPATH="${TORCH_MLIR_EXT_PYTHONPATH:-""}"
+pypath="${python_packages_dir}/torch_mlir"
 if [ ! -z ${TORCH_MLIR_EXT_PYTHONPATH} ]; then
-  PYTHONPATH="${PYTHONPATH}:${TORCH_MLIR_EXT_PYTHONPATH}"
+  pypath="${pypath}:${TORCH_MLIR_EXT_PYTHONPATH}"
 fi
 TORCH_MLIR_EXT_MODULES="${TORCH_MLIR_EXT_MODULES:-""}"
 ext_module="${ext_module:-""}"
@@ -37,10 +44,8 @@ fi
 set +u
 # To enable this python package, manually build torch_mlir with:
 #   -DTORCH_MLIR_ENABLE_PYTORCH_EXTENSIONS=ON
-#   -DTORCH_MLIR_ENABLE_JIT_IR_IMPORTER=ON
 # TODO: move this package out of JIT_IR_IMPORTER.
-echo $PYTHONPATH
-python3 \
+PYTHONPATH="${PYTHONPATH}:${pypath}" python3 \
   -m torch_mlir.jit_ir_importer.build_tools.torch_ods_gen \
   --torch_ir_include_dir="${torch_ir_include_dir}" \
   --pytorch_op_extensions="${ext_module}" \
diff --git a/build_tools/write_env_file.sh b/build_tools/write_env_file.sh
new file mode 100755
index 000000000000..8f3c9a59357f
--- /dev/null
+++ b/build_tools/write_env_file.sh
@@ -0,0 +1,26 @@
+#!/bin/bash
+# generate the .env file with default options.
+#
+# For arbitrary build/install directories, set the env variables:
+# - TORCH_MLIR_BUILD_DIR
+
+set -eu -o pipefail
+
+portable_realpath() {
+  # Create the directory if needed so that the `cd` doesn't fail.
+  mkdir -p "$1" && cd "$1" && pwd
+}
+
+td="$(portable_realpath "$(dirname "$0")"/..)"
+build_dir="$(portable_realpath "${TORCH_MLIR_BUILD_DIR:-$td/build}")"
+python_packages_dir="$build_dir/python_packages"
+
+write_env_file() {
+  echo "Updating $build_dir/.env file"
+  echo "PYTHONPATH=\"$(portable_realpath "$python_packages_dir/torch_mlir")\"" > "$build_dir/.env"
+  if ! cp "$build_dir/.env" "$td/.env"; then
+    echo "WARNING: Failed to write $td/.env"
+  fi
+}
+
+write_env_file
diff --git a/docs/development.md b/docs/development.md
index 372d4b9c4bf4..360dff8f9df8 100644
--- a/docs/development.md
+++ b/docs/development.md
@@ -197,14 +197,13 @@ TIP: add multiple target options to stack build phases
 #### Linux and macOS
 
 ```shell
-source $PWD/.env && export PYTHONPATH="${PYTHONPATH}:${PWD}/test/python/fx_importer"
+export PYTHONPATH=`pwd`/build/python_packages/torch_mlir:`pwd`/test/python/fx_importer
 ```
 
 #### Windows PowerShell
 
 ```shell
-$env:PYTHONPATH = ... # Get from .env file in torch-mlir repo root. 
-$env:PYTHONPATH += ";${PWD}/test/python/fx_importer"
+$env:PYTHONPATH = "$PWD/build/tools/torch-mlir/python_packages/torch_mlir;$PWD/test/python/fx_importer"
 ```
 
 ### Testing MLIR output in various dialects
@@ -215,6 +214,8 @@ Make sure you have activated the virtualenv and set the `PYTHONPATH` above
 (if running on Windows, modify the environment variable as shown above):
 
 ```shell
+source mlir_venv/bin/activate
+export PYTHONPATH=`pwd`/build/tools/torch-mlir/python_packages/torch_mlir:`pwd`/test/python/fx_importer
 python test/python/fx_importer/basic_test.py
 ```
 
@@ -225,10 +226,10 @@ using torchscript with the example `projects/pt1/examples/torchscript_resnet18_a
 This path doesn't give access to the current generation work that is being driven via the fx_importer
 and may lead to errors.
 
-Same as above, but you need to append to the PYTHONPATH:
+Same as above, but with different python path and example:
 
 ```shell
-export PYTHONPATH="${PYTHONPATH}:/projects/pt1/examples"
+export PYTHONPATH=`pwd`/build/tools/torch-mlir/python_packages/torch_mlir:`pwd`/projects/pt1/examples
 python projects/pt1/examples/torchscript_resnet18_all_output_types.py
 ```
 
@@ -258,6 +259,14 @@ jupyter notebook
 [Example IR](https://gist.github.com/silvasean/e74780f8a8a449339aac05c51e8b0caa) for a simple 1 layer MLP to show the compilation steps from TorchScript.
 
 
+### Interactive Use
+
+The `build_tools/write_env_file.sh` script will output a `.env`
+file in the workspace folder with the correct PYTHONPATH set. This allows
+tools like VSCode to work by default for debugging. This file can also be
+manually `source`'d in a shell.
+
+
 ### Bazel Build
 
 > **NOTE** Our Bazel build follows LLVM's Bazel build policy: only the