[pull] master from tensorflow:master - #8805
Merged
Merged
Conversation
Addresses #94130. Master already validates that grad's rank matches var's rank in the Sparse*Op family (landed separately, covers ResourceSparseApplyAdagradDA among others), but the existing regression test (testSparseApplyOpsRejectLowerRankGrad) only asserts on the exception type, not the message. Per gemini-code-assist review feedback, assertRaises alone could pass even if a different InvalidArgumentError fired for an unrelated reason. Added testResourceSparseApplyAdagradDAInvalidGradRank, which pins the specific "grad must have the same number of dimensions as var" message via assertRaisesRegex for this op, so the test can only pass if the intended validation actually fired.
CI on real GPU hardware failed: ResourceSparseApplyAdagradDA has no
registered GPU kernel (CPU only), so without explicit device placement
the op's variables default onto GPU:0 and the test fails with a device
placement error before ever reaching the validation code under test.
Wrapped the test body in `with ops.device("/cpu:0")` so it exercises the
CPU kernel deterministically regardless of what hardware CI runs on.
Match the TOTALORDER comparison attribute emitted by the current tf2xla legalization.
…ndexDomainIndex()`
For host-buffer transfers (such as `xla::ifrt::Client::CopyArraysToHostBufferShards`), the caller must identify the unique `IndexDomain`s for the array sharding and map them to corresponding shard indices. Currently, the user only has access to `xla::ifrt::Sharding{Spec}::IndexDomains()`, which returns one index domain per device shard without deduplication or native shard-to-unique-domain grouping. This is often expensive (`O(#devices)`) and is not in the form of the information that the user actually needs to use for the downstream API call.
This change adds new sharding APIs:
- `xla::ifrt::Sharding{Spec}::UniqueIndexDomains()`: Returns a vector of `IndexDomainAndShardIndices`, where each item contains a unique `IndexDomain` and an `absl::Span<const int>` view of shard indices mapped to that domain.
- `xla::ifrt::Sharding{Spec}::ShardToUniqueIndexDomainIndex()`: Returns an `absl::Span<const int>` mapping each global shard index to the index of its corresponding unique index domain.
The shard indices view and reverse mapping are cached independently inside the `ShardingSpec` instance and computed on-demand.
The unique index domains are always ordered deterministically in a manner agnostic of shapes and device lists, but specific to the sharding type for optimal traversal and consistent use across APIs.
PiperOrigin-RevId: 975419584
Implements `PjRtCApiClient::Compile(const XlaComputation&, CompileOptions)` so that clients can perform AoT compilation w/ the PJRT C API without immediately loading to device memory. PiperOrigin-RevId: 975423410
…n barrier [XLA][HLO Value Tracking] Update original value in AlgebraicSimplifier when pruning optimization barrier When AlgebraicSimplifier prunes unused elements from an optimization barrier, update the instruction original value for both the barrier and its new operand using CopyOriginalValue from hlo_original_value_util.h. Also add vector support to CopyOriginalValue for vector-based index mappings. PiperOrigin-RevId: 975441821
…pacting(alloc_constants, buffers_to_assign). - Add schedule-independent HloLiveRange liveness helpers (GetBuffersDefined, GetBuffersUsed, GetBytesDefined, GetParameterBytesAtStart, BufferLivesOut). - Add unit tests in hlo_live_range_test. PiperOrigin-RevId: 975443269
…ComputationOp across calls. PiperOrigin-RevId: 975451427
PiperOrigin-RevId: 975465130
PiperOrigin-RevId: 975511528
…nderflow PiperOrigin-RevId: 975522386
PiperOrigin-RevId: 975525604
…fle-comparator PiperOrigin-RevId: 975526150
PiperOrigin-RevId: 975526580
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
Sign up for free
to subscribe to this conversation on GitHub.
Already have an account?
Sign in.
10 participants
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.
See Commits and Changes for more details.
Created by
pull[bot] (v2.0.0-alpha.4)
Can you help keep this open source service alive? 💖 Please sponsor : )