Skip to content

Commit c2a5690

Browse files
hawkinspGoogle-ML-Automation
hawkinsp
authored and
Google-ML-Automation
committed
Port pretty-printer to C++.
This yields a significant speedup (3x) when printing large jaxprs. PiperOrigin-RevId: 768803336
1 parent 0d1b1ef commit c2a5690

18 files changed

+1265
-405
lines changed

‎build/requirements_lock_3_10.txt

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,10 +24,6 @@ cloudpickle==3.0.0 \
2424
--hash=sha256:246ee7d0c295602a036e86369c77fecda4ab17b506496730f2f576d9016fd9c7 \
2525
--hash=sha256:996d9a482c6fb4f33c1a35335cf8afd065d2a56e973270364840712d9131a882
2626
# via -r build/test-requirements.txt
27-
colorama==0.4.6 \
28-
--hash=sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44 \
29-
--hash=sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6
30-
# via -r build/test-requirements.txt
3127
contourpy==1.2.1 \
3228
--hash=sha256:00e5388f71c1a0610e6fe56b5c44ab7ba14165cdd6d695429c5cd94021e390b2 \
3329
--hash=sha256:10a37ae557aabf2509c79715cd20b62e4c7c28b8cd62dd7d99e5ed3ce28c3fd9 \

‎build/requirements_lock_3_11.txt

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,10 +24,6 @@ cloudpickle==3.0.0 \
2424
--hash=sha256:246ee7d0c295602a036e86369c77fecda4ab17b506496730f2f576d9016fd9c7 \
2525
--hash=sha256:996d9a482c6fb4f33c1a35335cf8afd065d2a56e973270364840712d9131a882
2626
# via -r build/test-requirements.txt
27-
colorama==0.4.6 \
28-
--hash=sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44 \
29-
--hash=sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6
30-
# via -r build/test-requirements.txt
3127
contourpy==1.2.1 \
3228
--hash=sha256:00e5388f71c1a0610e6fe56b5c44ab7ba14165cdd6d695429c5cd94021e390b2 \
3329
--hash=sha256:10a37ae557aabf2509c79715cd20b62e4c7c28b8cd62dd7d99e5ed3ce28c3fd9 \

‎build/requirements_lock_3_12.txt

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,10 +24,6 @@ cloudpickle==3.0.0 \
2424
--hash=sha256:246ee7d0c295602a036e86369c77fecda4ab17b506496730f2f576d9016fd9c7 \
2525
--hash=sha256:996d9a482c6fb4f33c1a35335cf8afd065d2a56e973270364840712d9131a882
2626
# via -r build/test-requirements.txt
27-
colorama==0.4.6 \
28-
--hash=sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44 \
29-
--hash=sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6
30-
# via -r build/test-requirements.txt
3127
contourpy==1.2.1 \
3228
--hash=sha256:00e5388f71c1a0610e6fe56b5c44ab7ba14165cdd6d695429c5cd94021e390b2 \
3329
--hash=sha256:10a37ae557aabf2509c79715cd20b62e4c7c28b8cd62dd7d99e5ed3ce28c3fd9 \

‎build/requirements_lock_3_13.txt

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,10 +24,6 @@ cloudpickle==3.0.0 \
2424
--hash=sha256:246ee7d0c295602a036e86369c77fecda4ab17b506496730f2f576d9016fd9c7 \
2525
--hash=sha256:996d9a482c6fb4f33c1a35335cf8afd065d2a56e973270364840712d9131a882
2626
# via -r build/test-requirements.txt
27-
colorama==0.4.6 \
28-
--hash=sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44 \
29-
--hash=sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6
30-
# via -r build/test-requirements.txt
3127
contourpy==1.3.0 \
3228
--hash=sha256:00ccd0dbaad6d804ab259820fa7cb0b8036bda0686ef844d24125d8287178ce0 \
3329
--hash=sha256:0be4d8425bfa755e0fd76ee1e019636ccc7c29f77a7c86b4328a9eb6a26d0639 \

‎build/requirements_lock_3_13_ft.txt

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,10 +24,6 @@ cloudpickle==3.1.0 \
2424
--hash=sha256:81a929b6e3c7335c863c771d673d105f02efdb89dfaba0c90495d1c64796601b \
2525
--hash=sha256:fe11acda67f61aaaec473e3afe030feb131d78a43461b718185363384f1ba12e
2626
# via -r build/test-requirements.txt
27-
colorama==0.4.6 \
28-
--hash=sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44 \
29-
--hash=sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6
30-
# via -r build/test-requirements.txt
3127
contourpy==1.3.1 \
3228
--hash=sha256:041b640d4ec01922083645a94bb3b2e777e6b626788f4095cf21abbe266413c1 \
3329
--hash=sha256:05e806338bfeaa006acbdeba0ad681a10be63b26e1b17317bfac3c5d98f36cda \

‎build/requirements_lock_3_14.txt

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,6 @@ build==1.2.2.post1
1010
# via -r build/test-requirements.txt
1111
cloudpickle==3.1.1
1212
# via -r build/test-requirements.txt
13-
colorama==0.4.6
14-
# via -r build/test-requirements.txt
1513
contourpy==1.3.2
1614
# via matplotlib
1715
cycler==0.12.1

‎build/test-requirements.txt

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
absl-py
22
cloudpickle
3-
colorama>=0.4.4
43
filelock
54
flatbuffers
65
hypothesis

‎jax/BUILD

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1155,7 +1155,8 @@ pytype_strict_library(
11551155
deps = [
11561156
":config",
11571157
":util",
1158-
] + py_deps("colorama"),
1158+
"//jax/_src/lib",
1159+
],
11591160
)
11601161

11611162
pytype_strict_library(

‎jax/_src/lib/__init__.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,12 @@ def _parse_version(v: str) -> tuple[int, ...]:
108108

109109
import jaxlib.weakref_lru_cache as weakref_lru_cache # noqa: F401
110110

111+
if jaxlib_extension_version >= 350:
112+
import jaxlib._pretty_printer as _pretty_printer # noqa: F401
113+
else:
114+
_pretty_printer = None
115+
116+
111117
# XLA garbage collection: see https://github.com/jax-ml/jax/issues/14882
112118
def _xla_gc_callback(*args):
113119
xla_client._xla.collect_garbage()

0 commit comments

Comments
 (0)