Yunlong Liu
Yunlong Liu
### Description Hi JAX team, We identify a bug with the JAX cuda plugin. Here is the writeup for the bug, https://docs.google.com/document/d/1ldlD8XQ6XYX4zcSRCUIVQyAUBJQZX6v9PdE2qX2_FGw/edit?usp=sharing To summarize, We accidentally found that an object...
A previous PR https://github.com/openxla/xla/pull/15170 adds a python binding for accessing the profiled instruction. Actually the API previously added contains repeated logic with `get_fdo_profile`, which works very similar as the added...
Hi, I found that even though `absl.testing` is extended from the unittest module but it is not compatible with `unittest disover` feature. This makes it hard to people who want...
### Description I have noticed that this commit completely turns off non-hermetic cuda toolkit in the build.py. https://github.com/google/jax/commit/599c13aa0964069fa21a9e164281581fefff3cd9 This doesn't seem to be great because in the open source community,...
Offline chatted with @hawkinsp on fixing the crash when using io_callback inside the custom partitioning. The test is currently skipped because of a JAX bug that the lifetime of the...
As titled. XLA/JAX both bumped bazel version fo 7.4.1. It is just a pain for the community to have multiple bazel versions locally.
``` W0000 00:00:1741812697.800457 69236 xplane_to_op_stats.cc:364] No CoreDetails found for TPU device plane: /device:GPU:0 2025-03-12 20:51:38.799314: F external/org_tensorflow/tensorflow/core/profiler/utils/kernel_stats_utils.cc:339] Check failed: stats.is_op_tensor_core_eligible == kernel_report.is_op_tensor_core_eligible() (1 vs. 0) Fatal Python error: Aborted Current...
Added loading for nvjitlink library version 12. On GB200, if we don't do this, sometimes we will get, ``` Traceback (most recent call last): RuntimeError: jaxlib/cuda/versions_helpers.cc:81: operation cusparseGetProperty(MAJOR_VERSION, &major) failed:...