I recently got stuck in using PyMC MMM tools due to Jax and Jaxlib dependency issues on Windows 11.

After searching a lot and not being able to find any straightforward solution for resolving the issue. The better solution seemed to be to work with WSL on Windows.

Here’s a quick setup guide for working with Python virtual environment in WSL through VS Code.
[Note: Attached blog post links in reference dive deeper into all the solutions available, use this just as a quick guide if you’re working with Windows 11 and are okay with Ubuntu distribution installation for WSL]

Setup WSL and Python

• Open Command Prompt or Windows PowerShell as administrator.
• Run the command wsl --install and restart the computer. This will install WSL along with Ubuntu Linux distribution.
• Python comes pre-installed in Ubuntu. Check if python is installed using python3 --version.
• Check the installation location of Python using which python
• Update all the packages using sudo apt update && sudo apt upgrade
• Update only Python 3 using sudo apt upgrade python3

Setting up pip and virtual environments

• Install pip using sudo apt install python3-pip
• venv is useful for managing virtual environments in python
• Install venv using sudo apt install python3-venv

Setup multiple python versions using deadsnakes ppa

• deadsnakes ppa contains multiple versions of python that can be installed and used in Ubuntu.
• Add deadsnakes ppa repository using sudo add-apt-repository ppa:deadsnakes/ppa

# Example installations
# install python 3.7
sudo apt install python 3.7-distutils
# install python 3.12
sudo apt install python 3.12-distutils

• After installation, multiple python versions can be used e.g., python3.12 --version command will use python 3.12 and python3.7 --version will use python 3.7

Setting up VS code for WSL

• Download and install VS Code (on Windows side)
• Install WSL extension in VS Code.

Complete the setup and launch your private virtual environment instance from WSL to VS code

• Create a virtual environment for your Python 3.12 version (choose the version you need) using python3.12 -m venv mmm_code. A folder named “mmm_code” should be created, change your folder name accordingly.
• Activate the virtual environment using the command source ./mmm_code/bin/activate
• The installed packages can be checked using the command python3 -m pip list
• Use code . to launch WSL instance in VS code (this will happen after completion of a few installations in WSL to support VS code).

Happy Coding!

References:

1. WSL Installation
2. Setup Ubuntu or Linux Distribution in Windows using WSL
3. Using WSL with Python in VS Code