Repository navigation
[pull] master from tensorflow:master - #8845
Merged
Merged
Conversation
In tensorflow/python/ops/linalg/linalg_impl.py, the Pade approximant implementations for matrix exponential (_matrix_exp_pade3, _matrix_exp_pade5, _matrix_exp_pade7, _matrix_exp_pade9, _matrix_exp_pade13) used unnormalized integer-cleared polynomial coefficients. For higher-degree approximants, the constant term b_0 scaled up to 1.76e10 (order 9) and 6.48e16 (order 13). Under nested forward-mode autodiff (ForwardAccumulator over ForwardAccumulator), evaluating tangents through matrix multiplications containing these unnormalized polynomial terms caused cross-terms to scale quadratically (O(b_k^2) ~ 10^20) and inflated matrix conditioning by 10 orders of magnitude. On a 2x2 rotation generator matrix at t = 0.7 where the exact analytical second derivative is 2*exp(0.7)*cos(1) ~ 2.176, nested forward autodiff returned 357,356,221. This commit normalizes all Pade polynomial coefficients by b_0 (c_k = b_k / b_0) following Nicholas J. Higham (2005) Equation (2.2), bounding all coefficients in [0, 1]. This eliminates condition-number inflation, preserves primal accuracy to machine precision, and prevents catastrophic derivative explosion in higher-order autodiff. Fixes #127224.
In tensorflow/python/ops/linalg/linalg_impl.py, drop the now-redundant c[0] * ident multiplication in _matrix_exp_pade3, _matrix_exp_pade5, _matrix_exp_pade7, _matrix_exp_pade9, and _matrix_exp_pade13. Since the previous commit normalizes coefficients by b_0, c[0] is always 1.0, so the multiplication was a no-op kernel call. In matrix_exponential_op_test.py, tighten the float64 tolerances in testSecondDerivativeNestedForwardAccumulator to rtol=atol=1e-12 for all three assertions, verifying the higher precision the normalization fix provides in float64.
Tightening testSecondDerivativeNestedForwardAccumulator's nested ForwardAccumulator (forward-over-forward) assertion to 1e-12 for float64 does not hold: isolating the matmul chain used by the Pade approximants (matrix_2, matrix_4, matrix_6, matrix_8) shows nested forward-mode autodiff diverges from the reverse-mode/analytical value by about 9% once three matmuls are chained, and the divergence is the same magnitude in float32 and float64. That rules out floating-point precision as the cause; it's a limitation in how TensorFlow composes forward-over-forward differentiation through chained MatMul ops, unrelated to the Pade coefficient conditioning this PR fixes. The 1st-order forward JVP and 2nd-order reverse-mode Hessian assertions keep the tightened 1e-12 tolerance for float64, since both verified to machine precision.
testSecondDerivativeNestedForwardAccumulator fails under graph mode's default control flow, where matrix_exponential's internal tf.cond (for the is_finite check) lowers to Switch/Merge ops that GradientTape/ForwardAccumulator can't differentiate through, raising NotImplementedError. Add test_util.run_v2_only so the test only runs under TF2 control flow, where forward- and reverse-mode differentiation through the op are supported.
…d arithmetic_optimizer_test PiperOrigin-RevId: 985048948
…sted-forward-ad PiperOrigin-RevId: 985061669
PiperOrigin-RevId: 985103347
Imported from GitHub PR openxla/xla#48931 📝 Summary of Changes Fix rbe pool assignment based on the tags, make sure lit testing tools get their solib dependencies into the runfiles. 🎯 Justification With the old implementation any cpu test which has gpu tag on it will be assigned to the single gpu pool. That is wrong we want cpu tests to be assigned to default pool so to not consume gpu runners. Additionally thi pr fixes the issue of executing lit tests under rbe while there are missing tools solib deps in their runfiles 🚀 Kind of Contribution Please remove what does not apply: ♻️ Cleanup 📊 Benchmark (for Performance Improvements) Not relevant 🧪 Unit Tests: CI 🧪 Execution Tests: CI Copybara import of the project: -- 7080bd83412cfa21b6a7d665a47952f4f7ecc09a by Alexandros Theodoridis <atheodor@amd.com>: Fix rbe pool assignment based on the tags, fix missing deps from lit tools -- 4dc1e1b09410b819f9be2db63593f3f3c501eb35 by Alexandros Theodoridis <atheodor@amd.com>: Address review comments -- b7e7a6b6463a003d9560dde5905ff4563a6088d9 by Alexandros Theodoridis <atheodor@amd.com>: Use gpu_fusion_to_mlir in gpu specific hlos -- 8f7ffa16c5a190d9efea763adc9c0522830fbdf1 by Alexandros Theodoridis <atheodor@amd.com>: Revert lit.bzl adaptions Merging this change closes #48931 PiperOrigin-RevId: 985128031
`KernelReuseCache` deduplicates kernel *compilation*, but every consumer deep copies the compiled binary out of the cache entry, because `CreateOwnedCubinCustomKernel` takes the CUBIN by value. A module that invokes the same kernel N times therefore keeps N copies of its CUBIN in memory, and the Triton fusion emitter copies it twice, because `TritonFusion::EmitResult` holds the cache entry by value. Add reference counted `SharedCudaCubinInMemory` and `SharedCudaPtxInMemory` payloads to `se::KernelLoaderSpec`, keep `KernelReuseCache::Entry::binary` in a `std::shared_ptr`, and switch the three cached emitter call sites to the new `CreateSharedCubinCustomKernel` factory. `KernelLoaderSpec` remains a value type and its accessors keep handing out non-owning views, so every reader - including `ToProto` - is unaffected. The two emitters that have no cache keep the owning factory, since they have nothing to share with. The serialized representation is unchanged; this only removes redundant copies in memory. Deduplicating the serialized form follows in a later change. PiperOrigin-RevId: 985139084
…wyu.py. `@com_google_absl//absl/strings` exports `string_view.h` in hdrs, so Bazel's layering_check allows including `absl/strings/string_view.h` when `@com_google_absl//absl/strings` is in deps. However, BANT has a hardcoded `absl_string_view_skip` special case that ignores `string_view.h` on `absl/strings:strings` and requires you to depend on `absl/strings:string_view`. This does not align with layering_check and the behaviour can't be configured or disabled. So this change filters out `@com_google_absl//absl/strings:string_view` findings (and paired `@com_google_absl//absl/strings` removal suggestions) from BANT's output in `run_dwyu.py` and which basically reverts this hard-coded rule in bant and allows the `absl/strings` target for depending on `strings/string_view.h` again. PiperOrigin-RevId: 985142787
When the scan dimension is tiled, the scan is emitted inside an `scf.for` loop
and the carry of each tile becomes the init of the next one. That carry never
compiled:
* `LowerScan::ExtractCarries` extracted the last element along the scan
dimension with a `tensor.extract_slice`, which has no lowering to the
TritonGPU dialect.
* For 1D scans, the rank-reduced carry and therefore the `scf.for` `iter_arg`
was a rank-zero tensor, which is illegal in TritonGPU
(`BlockedEncodingAttr::get` asserts rank >= 1).
Instead of rank-reducing the scan dimension, the inits and carries of
`xtile.scan` now keep it as a unit dimension and have the same rank as the scan
inputs and outputs. Carries are therefore never rank-zero, and the init is
broadcast over the tile with a `tt.broadcast`, both without any reshapes. A new
`xtile.scan` verifier enforces this.
`tt.scan` does not return the final carry, so it has to be extracted from the
scan output. Gathering it with a `tt.gather` would write the entire tile to
shared memory. Instead, mask all other elements to zero and select the
remaining one with an `or` reduction over the bitcast integer values, which is
exact because zero is the neutral element of `or`. This is the established
pattern for tiled Triton scans, see `triton_helpers.select_one` in PyTorch
inductor and proger/accelerated-scan. It reduces the shared memory of the added
tiled scan tests from 1024 to 16 bytes (1D) and from 4096 to 128 bytes (2D),
and lowers to a single `redux.sync.or.b32` per warp on sm_80 and later.
Also fix `GetDimsToSqueeze`, which keeps one unit dimension by popping the last
entry when all dimensions have size one. For a rank-zero tensor there is no
entry to pop and it read past the end of the `SmallVector`.
None of this was caught because the only scan test reaching the emitter tiled
the whole scan into a single tile, so the carry was dead and got eliminated.
Add tests that check the emitted IR and the numerics of tiled scans.
PiperOrigin-RevId: 985144827
Both XLA and TensorFlow have upgraded rules_cc to 0.2.20, which provides
the `@rules_cc//cc/libc:{musl,glibc}` config settings used by upstream
LLVM's `utils/bazel/llvm-project-overlay/llvm/config.bzl`. Remove the
obsolete patch to keep the number of patches to a minimum and prevent future
problems with patches failing to apply during LLVM integrates.
PiperOrigin-RevId: 985147212
PiperOrigin-RevId: 985178197
Imported from GitHub PR openxla/xla#49006 📝 Summary of Changes Fixes #49005. Fix `IsTopKIndexOutputUsed` so that returning either the complete TopK `(values, indices)` tuple or just its `indices` element as the computation result counts as using the indices. The check now recognizes a root TopK custom call and uses `!IsDead()` for the `get-tuple-element` instruction that extracts the indices. This preserves the values-only optimization for truly unused indices. 🎯 Justification A computation root is observable by its caller even when its instruction-user list is empty. Checking only instruction users misses these root uses and selects the values-only path, which substitutes a dummy index tensor with the global `{batch_size, k}` shape. For batch-sharded TopK, can produces a global-sized index output where a partition-local shape is required. The JAX reproducer observes `f32[8,50]` values and `s32[32,50]` indices on four TPU v4 devices, causing `jax.device_get` to fail. This change restores their correct partition-local shape, as verified by the HLO regression tests. 🚀 Kind of Contribution 🐛 Bug Fix, 🧪 Tests 🧪 Unit Tests: Added: - `TopKCustomCallTupleRootUsesIndices` - `TopKCustomCallIndexRootUsesIndices` Both tests shard a batch of 64 rows across two devices. They fail with the implementation change removed and pass with the fix: | HLO root case | Before | After | |---|---|---| | TopK tuple root | `(bf16[32,40], s32[64,40])` | `(bf16[32,40], s32[32,40])` | | Index get-tuple-element root | `s32[64,40]` | `s32[32,40]` | The following local test runs passed: ```bash bazel test //xla/service/spmd:spmd_partitioner_test \ --test_filter='*TopKCustomCall*RootUsesIndices*' bazel test //xla/service/spmd:spmd_partitioner_test \ --test_filter='*TopKCustomCall*' bazel test //xla/service/spmd:spmd_partitioner_test ``` The full test target passed all 10 shards, including the existing values-only TopK coverage. Formatting checks with clang-format 17.0.6 and `git diff --check` also passed. Copybara import of the project: -- 05dd889e8607699be735934bbd3a198d2354ec61 by yixiaoer <shin@yixiaoer.sg>: Fix SPMD TopK root index outputs Merging this change closes #49006 PiperOrigin-RevId: 985192875
We need to canonicalise fp8 case in advance. PiperOrigin-RevId: 985198861
Export the arena copy constructor and destructor for tensorflow::data::model::ModelProto_Node in _pywrap_tensorflow.def to fix undefined symbol linker errors on Windows. PiperOrigin-RevId: 985206485
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 : )