CI: 04/25/25 upstream sync #385
Open
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}
…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
When we run the program with "--xla_jf_bounds_check=true", we can selectively disable bounds checks for pallas kernels now. PiperOrigin-RevId: 748193719
PiperOrigin-RevId: 748220556
…pjit_test.py and array_test.py so we keep testing it until we fully delete it. PiperOrigin-RevId: 751083672
PiperOrigin-RevId: 751083687
… kernel. It's implemented as split dq and dkv kernels to have enough SMEM for double-buffering. The compute throughput is better than in the forwards pass, which is expected because there's less vector ops since we're not computing the softmax. PiperOrigin-RevId: 751084141
Reverts 492cd3d PiperOrigin-RevId: 751087903
And do other things... Co-authored-by: Yash Katariya <[email protected]>
PiperOrigin-RevId: 751103596
PiperOrigin-RevId: 751107289
PiperOrigin-RevId: 751109708
PiperOrigin-RevId: 751110300
PiperOrigin-RevId: 751112644
…g so we don't get the internal name for it in errors PiperOrigin-RevId: 751113771
PiperOrigin-RevId: 751118618
Co-authored-by: Yash Katariya <[email protected]>
…-shardings PiperOrigin-RevId: 751135041
PiperOrigin-RevId: 751139503
PiperOrigin-RevId: 751143638
PiperOrigin-RevId: 751148795
PiperOrigin-RevId: 751166905
…slice_p with Element block shapes PiperOrigin-RevId: 751173626
This is like the scan version of jax-ml#27970, which applied to while_loop. Fixes google/flax#4709
PiperOrigin-RevId: 751184927
…ization PiperOrigin-RevId: 751191568
Co-authored-by: Matthew Johnson <[email protected]> PiperOrigin-RevId: 751197878
PiperOrigin-RevId: 751212769
* Also add Python pipeline emitter support PiperOrigin-RevId: 751228049
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