# the latest tensorflow supports python 3.13
mm create -p ./env -y python=3.13
mm activate ./env
 
# add cuda lib paths to LD_LIBRARY_PATH automatically when activating the
# environment
mkdir -p $CONDA_PREFIX/etc/conda/activate.d
cat > $CONDA_PREFIX/etc/conda/activate.d/tensorflow-cuda.sh <<'EOF'
export _OLD_LD_LIBRARY_PATH="$LD_LIBRARY_PATH"
 
export LD_LIBRARY_PATH="$(python - <<'PY'
import site
from pathlib import Path
 
for sp in site.getsitepackages():
    nvidia = Path(sp) / "nvidia"
    if nvidia.exists():
        for lib in nvidia.glob("*/lib"):
            print(lib, end=":")
PY
)$LD_LIBRARY_PATH"
EOF
 
mkdir -p $CONDA_PREFIX/etc/conda/deactivate.d
cat > $CONDA_PREFIX/etc/conda/deactivate.d/tensorflow-cuda.sh <<'EOF'
export LD_LIBRARY_PATH="$_OLD_LD_LIBRARY_PATH"
EOF
 
pip install tensorflow[and-cuda]
 
. $CONDA_PREFIX/etc/conda/activate.d/tensorflow-cuda.sh
 
# make sure GPU is recognized
python -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"