mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fail gracefully in lobpcg_test if matplotlib isn't installed.

Browse Source 
There isn't yet a matplotlib that supports NumPy 2.0, so we need to support running tests without it.

PiperOrigin-RevId: 621228727
This commit is contained in:
Peter Hawkins 2024-04-02 10:47:55 -07:00 committed by jax authors
parent 00489be23d
commit 60458bb36a
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
tests/lobpcg_test.py
View File

@ -288,7 +288,10 @@ class LobpcgTest(jtu.JaxTestCase):
# We import matplotlib lazily because (a) it's faster this way, and
# (b) concurrent imports of matplotlib appear to trigger some sort of
# collision on the matplotlib cache lock on Windows.
from matplotlib import pyplot as plt
try:
from matplotlib import pyplot as plt
except (ModuleNotFoundError, ImportError):
return # If matplotlib isn't available, don't emit plots.
os.makedirs(lobpcg_debug_plot_dir, exist_ok=True)
clean_matrix_name = _clean_matrix_name(matrix_name)