Skip to content

Disable JAX GPU preallocation by default #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
YonghaoZhao722 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Disable JAX GPU preallocation by default #1

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d16b00c
Disable JAX GPU preallocation by default
YonghaoZhao722 May 8, 2026
3029a47
Merge pull request #1 from YonghaoZhao722/codex/disable-jax-gpu-memor…
YonghaoZhao722 May 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions mgw/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
from . import _jax_config # noqa: F401 - configure JAX before any OTT/JAX imports
from .models import PhiModel, PhiModelFFN, train_phi, train_phi_pro
from .geometry import pullback_metric_field, knn_graph, geodesic_distances, pairwise_squared_geodesic
from .gw import gw_distance, solve_gw
Expand Down
8 changes: 8 additions & 0 deletions mgw/_jax_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
"""Project-level JAX runtime defaults."""

import os

# Disable JAX GPU memory preallocation so OTT/JAX does not reserve most GPU
# memory in MGW workflows that also use PyTorch; set this env var before
# importing MGW to opt out and restore JAX's default allocator behavior.
os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
6 changes: 6 additions & 0 deletions mgw/gw.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,10 @@
import os

import numpy as np

# Match MGW package defaults even if this module is imported directly: avoid
# JAX/OTT preallocating most GPU memory alongside PyTorch workloads.
os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
from ott.problems.quadratic import quadratic_problem
from ott.solvers.linear.sinkhorn import Sinkhorn
from ott.solvers.quadratic.gromov_wasserstein import GromovWasserstein
Expand Down