CI: 04/16/25 upstream sync #367
Closed
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
![rocm-repo-management-api-2[bot]](https://avatars.githubusercontent.com/u/21157610?s=60&v=4){kind=small}
…antics), in addition to the existing ThreadSemantics (renamed to LoweringSemantics). UserThreadSemantics controls the thread semantics of the Pallas user's code, whereas LoweringSemantics controls the level at which Mosaic GPU emits code. PiperOrigin-RevId: 744857085
Also, add a job to the release test workflow that verifies that the release wheels can be installed. TESTED: 1. Full release: https://github.com/jax-ml/jax/actions/runs/14315832784 2. jax only release: https://github.com/jax-ml/jax/actions/runs/14316157252 PiperOrigin-RevId: 744857804
None is meant to represent the same thing as {replicated over all possible axes}. But without this canonicalization, we could compare None as not equal to {all possible axes}.
fixes jax-ml#26621
Unrelated: in several places, including the _check_rep path, we don't handle partial auto correctly, since we treat {all possible axes} as {all mesh axes}, but actually it should be more like {all mesh axes} - auto. We'll leave that fix for a follow-up...
PiperOrigin-RevId: 744868565
No internal code needs it any more. PiperOrigin-RevId: 744870756
…n the users wants to let SPMD decide the sharding. But this is a contradiction since layouts apply to device local shape and without knowing the sharding, you can't decide the layout. But there are cases where you don't care what the sharding is, you just want to force a row-major layout (for example). **This API should only be used for those cases**. PiperOrigin-RevId: 744888557
…uracy is not printed if it's None. PiperOrigin-RevId: 744889524
PiperOrigin-RevId: 744901130
PiperOrigin-RevId: 744905319
…xternal one is deprecated. PiperOrigin-RevId: 744908555
…mlir_module. PiperOrigin-RevId: 744934776
The ml_dtype requirement in JAX was updated to version 0.5.0+ (on Mar 20, 2025) - commit 4b7ead4 This update allows us to address previous FP8-related TODOs in jaxlib/XLA. PiperOrigin-RevId: 744943824
This enables exporting functions that use memory kinds to place data in different memories. jax-fixit PiperOrigin-RevId: 745008959
…ted test for parsing all canonical layouts. PiperOrigin-RevId: 745015431
PiperOrigin-RevId: 745038060
It is just an identity after the "stackless" rewrite. PiperOrigin-RevId: 745042532
If a function or class has a docstring, it does not need a `pass`. PiperOrigin-RevId: 745052107
The change in `utils.py` is to enable the use of `bitwidth` when the mlir dialect is not registered. PiperOrigin-RevId: 745060221
After this change we no longer skip tests that required 'RunState`. This necessitated a small fix in the pallas lowering of `while` and also enabling multiple i32 register bundling in the `optimization_barrier` lowering. PiperOrigin-RevId: 745065173
It defaults to True and is not flipped to False by any internal JAX users. PiperOrigin-RevId: 745067361
PiperOrigin-RevId: 745091128
…lf reference cycle loop in python. PiperOrigin-RevId: 745099538
PiperOrigin-RevId: 745119028
http://github.com/openxla/xla/commit/3764aee831189bd32a9c7dea56926b8f31ae86bf. PiperOrigin-RevId: 745130406
PiperOrigin-RevId: 745139008
PiperOrigin-RevId: 745145207
PiperOrigin-RevId: 747657069
…ptic. PiperOrigin-RevId: 747657234
PiperOrigin-RevId: 747663692
PiperOrigin-RevId: 747782436
PiperOrigin-RevId: 747848417
Not all of these have replacements, but I want to mark these as deprecated now so they can be removed in a future release. PiperOrigin-RevId: 747867544
http://github.com/openxla/xla/commit/0bfaa69a78c4306d6421f7ba78638427d9ce53e8. PiperOrigin-RevId: 747871119
PiperOrigin-RevId: 747889283
PiperOrigin-RevId: 747920008
PiperOrigin-RevId: 747920261
Shift ops expect RHS to have the same type as RHS, but a scalar was being used for RHS. PiperOrigin-RevId: 747934619
PiperOrigin-RevId: 747943230
Reverts b336daf PiperOrigin-RevId: 747988862
* A garbage collected type bound by nanobind is allocated with `PyType_GenericAlloc` (https://github.com/wjakob/nanobind/blob/8f245bd6e5544ff828beb46435af24b8769bfcdc/src/nb_type.cpp#L81). * `PyType_GenericAlloc` calls `_PyObject_GC_TRACK` (https://github.com/python/cpython/blob/884df116d79b05d9342e05e50484d61c684ecb8b/Objects/typeobject.c#L2357), and as soon as an object is tracked by the GC, it may be traversed (https://docs.python.org/3/c-api/gcsupport.html#c.PyObject_GC_Track). * However, at this point, the nanobind-bound object has been allocated, but its C++ constructor need not yet have run. The constructor runs when `__init__` is called. * In tp_traverse methods of nanobind-bound classes, we should test whether the instance is ready before visiting any C++ state. Instance readiness will be set at the end of the dispatch of `__init__` (https://github.com/wjakob/nanobind/blob/8f245bd6e5544ff828beb46435af24b8769bfcdc/src/nb_func.cpp#L857). In addition, all nanobind-bound classes have heap-allocated types. Make sure we visit `Py_TYPE(self)`. PiperOrigin-RevId: 747998575
PiperOrigin-RevId: 748027325
…led before. PiperOrigin-RevId: 748027360
PiperOrigin-RevId: 748059518
PiperOrigin-RevId: 748059520
…Array` data Generally, we want to maintain that key data backing a `PRNGKeyArray` is a `jax.Array`. This change converts NumPy arrays on construction. Co-authored-by: Yash Katariya <[email protected]> PiperOrigin-RevId: 748077900
PiperOrigin-RevId: 748085747
The first letter was inadvertently made lower-case in the previous re-naming CL. PiperOrigin-RevId: 748086763
PiperOrigin-RevId: 748087650
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Daily sync with upstream