Use convert_to_tensor in __init__ to avoid eager-only .numpy()#1
Open
CodersAcademy006 wants to merge 1820 commits intomasterfrom
Open
Use convert_to_tensor in __init__ to avoid eager-only .numpy()#1CodersAcademy006 wants to merge 1820 commits intomasterfrom
CodersAcademy006 wants to merge 1820 commits intomasterfrom
Conversation
PiperOrigin-RevId: 834642693
…sm_110) Imported from GitHub PR openxla/xla#34107 This fixes the failing test StatelessAutotunerTest.CublasFallbackForBf16Bf16F32Algorithm on Jetson Thor GPUs with compute capability 11.0. Copybara import of the project: -- 2385bc08e0a8778d3acfa48caa4a885dc259b206 by Andrei Ivanov <anivanov@nvidia.com>: [XLA:GPU] Fix cublas fallback test on Thor GPU (sm_110) Merging this change closes tensorflow#34107 PiperOrigin-RevId: 834674818
PiperOrigin-RevId: 834674902
Imported from GitHub PR openxla/xla#34103 📝 Summary of Changes "mandatory" compatible layouts have to be assigned to both operands and outputs simultaneously such that subsequent layout propagation does not alter one of them making the operation invalid. 🚀 Kind of Contribution 🐛 Bug Fix 🧪 Unit Tests: yes 🧪 Execution Tests: no Copybara import of the project: -- f0ff62e4bf031a3aebf4cdadb66139b3b1120307 by Ilia Sergachev <isergachev@nvidia.com>: [GPU] Fix layout assignment of bitcast-converts. "mandatory" compatible layouts have to be assigned to both operands and outputs simultaneously such that subsequent layout propagation does not alter one of them making the operation invalid. Merging this change closes tensorflow#34103 PiperOrigin-RevId: 834676734
…intrinsics. PiperOrigin-RevId: 834684558
…parameter``` is true. Old GSPMD propagation needs them since they do not have the concept of open/closed sharding. In Shardy with sdy-round-trip, JAX creates the correct open/closed shardings for parameters and results. We do not need these vectors at all. Before this change, we always canonicalize layout at after Shardy propagaiton, which may be redundant. PiperOrigin-RevId: 834684933
PiperOrigin-RevId: 834691273
…tions Imported from GitHub PR openxla/xla#32053 📝 Summary of Changes Added CommandBuffer support for Convolution ops graph capture of convolutions is enabled only for convolution custom call targets explicitly added to '--legacy_command_buffer_custom_call_targets' list: see command_buffer_scheduling_test.cc for an example. 🎯 Justification This op wase missing for whatever reason: this results in graph fragmentation especially for large models. Hence one gets several (sometimes many) execution graphs instead of just one. 🚀 Kind of Contribution ✨ New Feature 🧪 Unit Tests: Added new subtest to xla/service/gpu/transforms/command_buffer_scheduling_test.cc This is a splitted PR originated from openxla/xla#30855 @xla-rotation could you have a look please ? Copybara import of the project: -- 5f5f5bc8ba8212ceb6afde6f9729ba4a951e4051 by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>: adding coll permute and convolution to command buffers added UTs and convolution command test fixes added rebase fixes capture only those convolution targets which are explictly Revert "adding coll permute and convolution to command buffers" This reverts commit 75847e67261b4589162411c9846ed9c0b9fc1ed5. added conv to command buffers fixing build and test -- e8afa3296a4a8ad079cde2e84391c7e0006ddf52 by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>: fixing build -- b529288221708e59a04bfaacdfa5b7a1c25b091e by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>: rewritten ConvolutionCmd, adapted command_buffer_conv_pass -- 3ecd3d0516a7766d70d636b2110b4b310a9be7b2 by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>: some cosmetics Merging this change closes tensorflow#32053 PiperOrigin-RevId: 834708288
PiperOrigin-RevId: 834762307
PiperOrigin-RevId: 834776465
…numbers from a structured MLIR attribute. The `ParseDimensionNumbers` function now expects a single `dimension_numbers` attribute in the composite attributes, structured as an array of arrays representing contracting and batch dimensions. This change simplifies the attribute parsing. Additionally, checks are added to ensure that the parsed dimension numbers are supported and to handle cases where scale factors are not divisible by 32 for non-BF16 types, preventing rewriting in those scenarios. The test case is updated to reflect the new attribute format. PiperOrigin-RevId: 834777398
…rror flakes. They may be causing permission errors on Windows when Bazel tries to access header files in Windows SDK/Clang while building @@bazel_tools targets. PiperOrigin-RevId: 834792776
In LiteRT, input and output names shouldn't be empty. Populate default names if tensors don't have names. PiperOrigin-RevId: 834808503
PiperOrigin-RevId: 834837889
PiperOrigin-RevId: 834839651
…ation model for collective-permute uses transfer size and communication pattern type as input key. The interpolation is implemented as follows: - `CollectivePermuteCostModelType` which classifies communication patterns (e.g. one-way, two-way-all-mutual) is added to `ExactInterpolatorKey` for collective-permute instructions. - Interpolation for collective-permute uses only exact matching via `ExactInterpolator` based on transfer size, and does not use `FallbackInterpolator`. This is because cost of collective-permute is primarily dependent on bytes transferred and communication pattern, and not on number of devices in the same way as other collectives. PiperOrigin-RevId: 834846468
PiperOrigin-RevId: 834846503
Updates LLVM usage to match [355e0f94af5a](llvm/llvm-project@355e0f94af5a) PiperOrigin-RevId: 834865231
…rvice New namespace is absl_testing:: PiperOrigin-RevId: 834872277
When the call splitter is called on a non-flat graph, we don't want to implicitly flatten it by creating new bodies at each callsite. PiperOrigin-RevId: 834875376
PiperOrigin-RevId: 834896529
PiperOrigin-RevId: 834972320
Fixes a problem where structured_tensor fails to import under Python 3.14 with the error:
```
File ".../tensorflow/python/ops/structured/structured_tensor.py", line 54, in <module>
class StructuredTensor(extension_type.BatchableExtensionType):
...<1135 lines>...
return self._ragged_shape.rank
File ".../tensorflow/python/framework/extension_type.py", line 90, in __init__
_check_field_annotations(cls)
~~~~~~~~~~~~~~~~~~~~~~~~^^^^^
File ".../tensorflow/python/framework/extension_type.py", line 935, in _check_field_annotations
raise ValueError(
...<2 lines>...
)
ValueError: The field annotations for StructuredTensor are invalid. Field FieldName is missing a type annotation.
```
PiperOrigin-RevId: 834982237
PiperOrigin-RevId: 834982495
PiperOrigin-RevId: 834997967
…host transfer extension. PiperOrigin-RevId: 835038444
PiperOrigin-RevId: 835066344
PiperOrigin-RevId: 835085920
PiperOrigin-RevId: 835096840
PiperOrigin-RevId: 835111311
…permutation, which operate directly on containers. This simplifies the code by removing the need to pass begin() and end() iterators. Also includes minor refactoring of if/else structures and updates function signatures for clarity. PiperOrigin-RevId: 837296580
PiperOrigin-RevId: 837297785
PiperOrigin-RevId: 837301178
…ABSL versions internally and in OSS. PiperOrigin-RevId: 837314334
…already writes the chosen node into `ready_chosen` using `memcpy`, so we don't need the conditionals in the print statement. Also, we need to save the "unchosen" node's information for printing purposes. PiperOrigin-RevId: 837321312
To avoid bringing GPU dependencies with default FFI target, split GPU-specific context decoding into backends/gpu:ffi PiperOrigin-RevId: 837342185
This is expected to be a safe change. Calls to multimap::find() historically returned the first matching element, or end() if there is no match. However, this is not guaranteed, and recent changes in libc++ changed this to return an arbitrary element that matches. Using equal_range() is a safe replacement that will preserve the current behavior. PiperOrigin-RevId: 837348099
… id mapping PiperOrigin-RevId: 837366466
PiperOrigin-RevId: 837373785
PiperOrigin-RevId: 837386345
…s one was missed in previous consolidations. PiperOrigin-RevId: 837391515
PiperOrigin-RevId: 837393069
…tx backend Instead of relying on is_autotuning_compilation boolean it's now up to backend runner to properly set --fail_ptx_compilation_on_register_spilling based on --xla_gpu_filter_kernels_spilling_registers_on_autotuning. Note the change in AutotunerCompileUtil::Compile calls GpuCodegenBackend::AdjustDebugOptionsForAutotuning. That seems to also improve compile time of benchmarks. Dropped gemm fusion autotuner tests about spilling as they basically tested if the backend respects debug option flags. With this change they would become tautological at best and don't exercise any behavior of gemm_fusion_autotuner. PiperOrigin-RevId: 837403708
…nalysis. PiperOrigin-RevId: 837407555
Prior to this change we would hit an assert when constructing the tensor::CollapseShapeOp in the 0D->0D case. PiperOrigin-RevId: 837419131
Can be enabled with XLA_FLAGS="--xla_backend_extra_options=xla_cpu_enable_tiled_emitter" (!warning! may not work as expected for now) Reverts 3a4bd44 PiperOrigin-RevId: 837419373
…er base class. PiperOrigin-RevId: 837424830
PiperOrigin-RevId: 837425105
PiperOrigin-RevId: 837425159
…uctions to tiled emitter. PiperOrigin-RevId: 837432186
PiperOrigin-RevId: 837446324
flag Corresponding field is already removed from the proto PiperOrigin-RevId: 837451384
PiperOrigin-RevId: 837453532
To get the behavior in line with internal ASSERT_OK_AND_ASSIGN. Unlike the non-TF_ counterpart, the macro expands to code that already includes the trailing semicolon so it happens to work. Attempts to replace it with an internal variant as part of b/444419873 make it not work anymore, until the missing semicolons are added. PiperOrigin-RevId: 837461910
Some arguments to kernels may not be managed by the buffer assignments and consequently have no buffer slices attached to them. Examples of these include scalars and arguments whose memory is managed by the runtime thunks [CollectiveKernelThunk]. This change introduces a new type for arguments to be passed in with a shape but without an associated slice. PiperOrigin-RevId: 837470085
This change replaces usages of tsl::errors::AlreadyExists with absl::AlreadyExistsError, wrapping arguments in absl::StrCat where necessary. This addresses deprecation warnings and moves towards standard Abseil error handling. Changes: - Replaced errors::AlreadyExists with absl::AlreadyExistsError. - Used absl::StrCat to construct error messages where necessary. Reverts 2fc3b48 PiperOrigin-RevId: 837473556
…o_tensor(...) to avoid .numpy() in graph/XLA mode
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 join this conversation on GitHub.
Already have an account?
Sign in to comment
20 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.
Replace tf.constant(...) calls inside init examples/tests with tf.convert_to_tensor(...) so classes work under @tf.function(jit_compile=True). This change is conservative and only updates examples and the test module. Please review.