ai-edge-tensorflow 2.21.0.dev20250805__cp39-cp39-macosx_12_0_arm64.whl → 2.21.0.dev20250818__cp39-cp39-macosx_12_0_arm64.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.

This version of ai-edge-tensorflow might be problematic. Click here for more details.

Files changed (455) hide show
  1. {ai_edge_tensorflow-2.21.0.dev20250805.dist-info → ai_edge_tensorflow-2.21.0.dev20250818.dist-info}/METADATA +2 -2
  2. {ai_edge_tensorflow-2.21.0.dev20250805.dist-info → ai_edge_tensorflow-2.21.0.dev20250818.dist-info}/RECORD +451 -422
  3. tensorflow/_api/v2/__internal__/distribute/multi_process_runner/__init__.py +6 -6
  4. tensorflow/_api/v2/compat/v1/raw_ops/__init__.py +13 -13
  5. tensorflow/_api/v2/compat/v2/__internal__/distribute/multi_process_runner/__init__.py +6 -6
  6. tensorflow/_api/v2/compat/v2/raw_ops/__init__.py +13 -13
  7. tensorflow/_api/v2/raw_ops/__init__.py +13 -13
  8. tensorflow/compiler/mlir/stablehlo/stablehlo_extension.so +0 -0
  9. tensorflow/compiler/tf2xla/ops/_xla_ops.so +0 -0
  10. tensorflow/compiler/xla/service/hlo_pb2.py +80 -80
  11. tensorflow/compiler/xla/xla_data_pb2.py +89 -93
  12. tensorflow/core/protobuf/config_pb2.py +27 -27
  13. tensorflow/include/external/KleidiAI/kai/kai_common.h +3 -12
  14. tensorflow/include/external/KleidiAI/kai/ukernels/matmul/matmul_clamp_bf16_qai8dxp_qsi4cxp/kai_matmul_clamp_bf16_qai8dxp1x8_qsi4cxp8x8_1x8_neon_dotprod.h +139 -0
  15. tensorflow/include/external/KleidiAI/kai/ukernels/matmul/matmul_clamp_bf16_qai8dxp_qsi4cxp/kai_matmul_clamp_bf16_qai8dxp4x8_qsi4cxp8x8_8x8_neon_i8mm.h +139 -0
  16. tensorflow/include/external/KleidiAI/kai/ukernels/matmul/matmul_clamp_bf16_qai8dxp_qsi4cxp/kai_matmul_clamp_bf16_qai8dxp_qsi4cxp_interface.h +52 -0
  17. tensorflow/include/external/KleidiAI/kai/ukernels/matmul/matmul_clamp_f16_f16_f16p/kai_matmul_clamp_f16_f16_f16p2vlx2b_1x8vl_sme_mla.h +119 -0
  18. tensorflow/include/external/KleidiAI/kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p2vlx1b_1x8vl_sme_mla.h +114 -0
  19. tensorflow/include/external/KleidiAI/kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p_interface.h +1 -1
  20. tensorflow/include/external/KleidiAI/kai/ukernels/matmul/pack/kai_lhs_quant_pack_qai8dxp_bf16_neon.h +77 -0
  21. tensorflow/include/external/KleidiAI/kai/ukernels/matmul/pack/kai_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme.h +4 -4
  22. tensorflow/include/external/XNNPACK/include/experimental.h +34 -7
  23. tensorflow/include/external/llvm-project/mlir/include/mlir/Analysis/DataFlowFramework.h +31 -23
  24. tensorflow/include/external/llvm-project/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h +8 -0
  25. tensorflow/include/external/llvm-project/mlir/include/mlir/Conversion/Passes.h.inc +43 -6
  26. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.h.inc +2 -1
  27. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/ArmNeon/ArmNeon.cpp.inc +4 -4
  28. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/ArmSVE/IR/ArmSVE.cpp.inc +5 -5
  29. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/Async/IR/AsyncOps.h.inc +3 -3
  30. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h +14 -13
  31. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/CommonFolders.h +114 -33
  32. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/DLTI/DLTIAttrs.cpp.inc +4 -4
  33. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/Func/IR/FuncOps.h.inc +3 -3
  34. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/GPU/IR/GPUOps.cpp.inc +443 -190
  35. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/GPU/IR/GPUOps.h.inc +177 -97
  36. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/GPU/IR/GPUOpsAttributes.cpp.inc +3 -3
  37. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicConversions.inc +37 -2
  38. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicFromLLVMIRConversions.inc +1310 -1480
  39. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.cpp.inc +1096 -397
  40. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.h.inc +527 -257
  41. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.cpp.inc +42 -0
  42. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.h.inc +4 -4
  43. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/LLVMIR/LLVMOpsAttrDefs.cpp.inc +4 -4
  44. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/LLVMIR/LLVMOpsDialect.h.inc +3 -0
  45. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.cpp.inc +2 -2
  46. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/LLVMIR/NVVMConversions.inc +15 -0
  47. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.cpp.inc +644 -118
  48. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.h.inc +354 -18
  49. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/LLVMIR/NVVMOpsAttributes.cpp.inc +216 -0
  50. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/LLVMIR/NVVMOpsAttributes.h.inc +41 -0
  51. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/LLVMIR/NVVMOpsEnums.cpp.inc +33 -0
  52. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/LLVMIR/NVVMOpsEnums.h.inc +104 -0
  53. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/LLVMIR/ROCDLConversions.inc +8 -0
  54. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.cpp.inc +213 -0
  55. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.h.inc +169 -0
  56. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/Linalg/IR/Linalg.h +190 -4
  57. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h.inc +2 -2
  58. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yamlgen.cpp.inc +0 -414
  59. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yamlgen.td +0 -322
  60. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.h.inc +22 -3
  61. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc +4668 -6104
  62. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.h.inc +2723 -3810
  63. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/Linalg/Passes.capi.cpp.inc +8 -0
  64. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/Linalg/Passes.capi.h.inc +5 -0
  65. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/Linalg/Passes.h.inc +176 -0
  66. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h +11 -0
  67. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/MLProgram/IR/MLProgramOps.h.inc +2 -2
  68. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.cpp.inc +383 -296
  69. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.h.inc +58 -18
  70. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/OpenMP/OpenMPOpsAttributes.cpp.inc +165 -0
  71. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/OpenMP/OpenMPOpsAttributes.h.inc +21 -0
  72. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/PDLInterp/IR/PDLInterpOps.h.inc +1 -1
  73. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/Quant/IR/QuantTypes.h +2 -0
  74. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/SCF/IR/SCFOps.cpp.inc +173 -21
  75. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/SCF/IR/SCFOps.h.inc +81 -6
  76. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/SCF/Transforms/Passes.h +3 -0
  77. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/SCF/Transforms/Passes.h.inc +130 -0
  78. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/SCF/Transforms/Transforms.h +6 -0
  79. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVAttributes.cpp.inc +3 -3
  80. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVEnumAvailability.cpp.inc +1 -0
  81. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVEnums.cpp.inc +12 -0
  82. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVEnums.h.inc +11 -7
  83. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVOpAvailabilityImpl.inc +56 -0
  84. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVOps.cpp.inc +711 -159
  85. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVOps.h.inc +425 -42
  86. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVSerialization.inc +30 -0
  87. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVTypes.h +38 -2
  88. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.h.inc +1 -1
  89. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.cpp.inc +2 -2
  90. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorTypes.cpp.inc +2 -2
  91. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/Transform/IR/TransformOps.h.inc +2 -2
  92. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/Transform/IR/TransformTypes.cpp.inc +1 -1
  93. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/Vector/IR/VectorOps.cpp.inc +1394 -2072
  94. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/Vector/IR/VectorOps.h +1 -0
  95. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/Vector/IR/VectorOps.h.inc +183 -474
  96. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h +1 -1
  97. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/X86Vector/X86Vector.h.inc +2 -2
  98. tensorflow/include/external/llvm-project/mlir/include/mlir/IR/BuiltinAttributes.cpp.inc +2 -2
  99. tensorflow/include/external/llvm-project/mlir/include/mlir/IR/Diagnostics.h +4 -0
  100. tensorflow/include/external/llvm-project/mlir/include/mlir/IR/Operation.h +20 -0
  101. tensorflow/include/external/llvm-project/mlir/include/mlir/IR/PatternMatch.h +14 -0
  102. tensorflow/include/external/llvm-project/mlir/include/mlir/IR/StorageUniquerSupport.h +1 -1
  103. tensorflow/include/external/llvm-project/mlir/include/mlir/IR/SymbolInterfaces.h.inc +14 -14
  104. tensorflow/include/external/llvm-project/mlir/include/mlir/Interfaces/CallInterfaces.cpp.inc +52 -12
  105. tensorflow/include/external/llvm-project/mlir/include/mlir/Interfaces/CallInterfaces.h.inc +178 -136
  106. tensorflow/include/external/llvm-project/mlir/include/mlir/Interfaces/FunctionInterfaces.cpp.inc +25 -25
  107. tensorflow/include/external/llvm-project/mlir/include/mlir/Interfaces/FunctionInterfaces.h.inc +28 -16
  108. tensorflow/include/external/llvm-project/mlir/include/mlir/Pass/PassOptions.h +1 -1
  109. tensorflow/include/external/llvm-project/mlir/include/mlir/Support/ToolUtilities.h +15 -0
  110. tensorflow/include/external/llvm-project/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h +18 -5
  111. tensorflow/include/external/llvm-project/mlir/include/mlir/Transforms/DialectConversion.h +3 -0
  112. tensorflow/include/external/llvm-project/mlir/lib/AsmParser/Lexer.h +3 -0
  113. tensorflow/include/external/shardy/shardy/dialect/sdy/ir/ops.h.inc +22 -21
  114. tensorflow/include/external/shardy/shardy/dialect/sdy/transforms/export/explicit_reshards_util.h +16 -9
  115. tensorflow/include/external/shardy/shardy/dialect/sdy/transforms/export/passes.h.inc +21 -143
  116. tensorflow/include/external/stablehlo/stablehlo/dialect/Base.h +5 -0
  117. tensorflow/include/external/stablehlo/stablehlo/dialect/StablehloAttrs.cpp.inc +2 -2
  118. tensorflow/include/external/stablehlo/stablehlo/dialect/StablehloOps.cpp.inc +48 -28
  119. tensorflow/include/external/stablehlo/stablehlo/dialect/StablehloOps.h +13 -0
  120. tensorflow/include/external/stablehlo/stablehlo/dialect/StablehloOps.h.inc +1 -0
  121. tensorflow/include/external/stablehlo/stablehlo/dialect/Version.h +1 -1
  122. tensorflow/include/external/stablehlo/stablehlo/dialect/VhloAttrs.cpp.inc +6 -6
  123. tensorflow/include/external/stablehlo/stablehlo/dialect/VhloTypeDefs.cpp.inc +123 -8
  124. tensorflow/include/external/stablehlo/stablehlo/dialect/VhloTypeDefs.h.inc +30 -0
  125. tensorflow/include/tensorflow/compiler/jit/flags.h +1 -0
  126. tensorflow/include/tensorflow/compiler/jit/tf_graph_to_hlo_compiler.h +4 -3
  127. tensorflow/include/tensorflow/compiler/jit/xla_compiler_options_util.h +3 -0
  128. tensorflow/include/tensorflow/compiler/mlir/tensorflow/ir/tf_all_ops.h.inc +334 -12
  129. tensorflow/include/tensorflow/compiler/mlir/tensorflow/ir/tf_device.h.inc +1 -1
  130. tensorflow/include/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.h.inc +1 -1
  131. tensorflow/include/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.h.inc +333 -11
  132. tensorflow/include/tensorflow/compiler/mlir/tensorflow/transforms/passes.h +3 -0
  133. tensorflow/include/tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc +130 -0
  134. tensorflow/include/tensorflow/compiler/mlir/tensorflow/utils/deserialize_mlir_module_utils.h +34 -0
  135. tensorflow/include/tensorflow/compiler/mlir/tensorflow/utils/serialize_mlir_module_utils.h +0 -10
  136. tensorflow/include/tensorflow/compiler/tf2xla/sharding_util.h +11 -0
  137. tensorflow/include/tensorflow/compiler/tf2xla/xla_compiled_cpu_function.h +1 -10
  138. tensorflow/include/tensorflow/compiler/tf2xla/xla_compiler.h +27 -3
  139. tensorflow/include/tensorflow/compiler/xla/backends/cpu/codegen/emitters/cpu_scatter_emitter.h +4 -3
  140. tensorflow/include/tensorflow/compiler/xla/backends/cpu/codegen/execution_engine.h +1 -1
  141. tensorflow/include/tensorflow/compiler/xla/backends/cpu/codegen/fusion_compiler.h +11 -5
  142. tensorflow/include/tensorflow/compiler/xla/backends/cpu/codegen/ir_compiler.h +1 -0
  143. tensorflow/include/tensorflow/compiler/xla/backends/cpu/{onednn_fusion.h → onednn_support.h} +4 -4
  144. tensorflow/include/tensorflow/compiler/xla/backends/cpu/runtime/thunk.h +108 -58
  145. tensorflow/include/tensorflow/compiler/xla/backends/cpu/runtime/thunk_executor.h +11 -20
  146. tensorflow/include/tensorflow/compiler/xla/backends/cpu/runtime/work_queue.h +3 -13
  147. tensorflow/include/tensorflow/compiler/xla/backends/cpu/runtime/xnnpack/xnn_convolution_thunk.h +1 -1
  148. tensorflow/include/tensorflow/compiler/xla/backends/cpu/runtime/xnnpack/xnn_dot_thunk.h +1 -1
  149. tensorflow/include/tensorflow/compiler/xla/backends/cpu/runtime/xnnpack/xnn_fusion_thunk.h +20 -19
  150. tensorflow/include/tensorflow/compiler/xla/backends/cpu/runtime/xnnpack/xnn_interop.h +45 -0
  151. tensorflow/include/tensorflow/compiler/xla/backends/cpu/runtime/xnnpack/xnn_threadpool.h +11 -27
  152. tensorflow/include/tensorflow/compiler/xla/backends/cpu/transforms/onednn_matcher.h +1 -1
  153. tensorflow/include/tensorflow/compiler/xla/backends/cpu/transforms/xnn_matcher.h +5 -5
  154. tensorflow/include/tensorflow/compiler/xla/backends/cpu/xnn_emitter.h +2 -2
  155. tensorflow/include/tensorflow/compiler/xla/backends/cpu/xnn_gemm_config.h +1 -0
  156. tensorflow/include/tensorflow/compiler/xla/backends/cpu/{xnn_fusion.h → xnn_support.h} +7 -8
  157. tensorflow/include/tensorflow/compiler/xla/backends/gpu/codegen/emitters/ir/_virtual_includes/xla_gpu_ops_inc_gen/xla_gpu_ops.h.inc +2 -2
  158. tensorflow/include/tensorflow/compiler/xla/backends/gpu/codegen/emitters/ir/xla_gpu_ops.h.inc +2 -2
  159. tensorflow/include/tensorflow/compiler/xla/client/executable_build_options.h +5 -0
  160. tensorflow/include/tensorflow/compiler/xla/codegen/emitters/concatenate_kernel_emitter.h +0 -1
  161. tensorflow/include/tensorflow/compiler/xla/codegen/emitters/ir/_virtual_includes/xla_ops_inc_gen/xla_ops.h.inc +1 -1
  162. tensorflow/include/tensorflow/compiler/xla/codegen/emitters/ir/xla_ops.h.inc +1 -1
  163. tensorflow/include/tensorflow/compiler/xla/codegen/emitters/loop_kernel_emitter.h +0 -1
  164. tensorflow/include/tensorflow/compiler/xla/codegen/emitters/transforms/pass_pipelines.h +28 -0
  165. tensorflow/include/tensorflow/compiler/xla/codegen/emitters/transforms/passes.h +1 -1
  166. tensorflow/include/tensorflow/compiler/xla/codegen/emitters/transforms/passes.h.inc +92 -92
  167. tensorflow/include/tensorflow/compiler/xla/codegen/emitters/utils.h +3 -1
  168. tensorflow/include/{xla/codegen/math → tensorflow/compiler/xla/codegen/intrinsic}/erf.h +4 -4
  169. tensorflow/include/{xla/codegen/math → tensorflow/compiler/xla/codegen/intrinsic}/exp.h +4 -4
  170. tensorflow/include/{xla/codegen/math → tensorflow/compiler/xla/codegen/intrinsic}/fptrunc.h +4 -4
  171. tensorflow/include/tensorflow/compiler/xla/codegen/{math → intrinsic}/intrinsic.h +29 -3
  172. tensorflow/include/tensorflow/compiler/xla/codegen/{math/math_compiler_lib.h → intrinsic/intrinsic_compiler_lib.h} +5 -5
  173. tensorflow/include/tensorflow/compiler/xla/codegen/{math → intrinsic}/ldexp.h +4 -4
  174. tensorflow/include/tensorflow/compiler/xla/codegen/{math → intrinsic}/log1p.h +4 -4
  175. tensorflow/include/{xla/codegen/math → tensorflow/compiler/xla/codegen/intrinsic}/rsqrt.h +9 -9
  176. tensorflow/include/{xla/codegen/math → tensorflow/compiler/xla/codegen/intrinsic}/string_interner.h +5 -5
  177. tensorflow/include/{xla/codegen/math → tensorflow/compiler/xla/codegen/intrinsic}/tanh.h +6 -9
  178. tensorflow/include/{xla/codegen/math → tensorflow/compiler/xla/codegen/intrinsic}/vec_name_mangler.h +5 -5
  179. tensorflow/include/{xla/codegen/math_lib.h → tensorflow/compiler/xla/codegen/intrinsic_lib.h} +16 -18
  180. tensorflow/include/tensorflow/compiler/xla/codegen/trace_pass_instrumentation.h +47 -0
  181. tensorflow/include/tensorflow/compiler/xla/core/collectives/communicator.h +10 -0
  182. tensorflow/include/tensorflow/compiler/xla/fp_util.h +3 -1
  183. tensorflow/include/tensorflow/compiler/xla/hlo/analysis/indexing_analysis.h +14 -2
  184. tensorflow/include/tensorflow/compiler/xla/hlo/builder/xla_builder.h +14 -0
  185. tensorflow/include/tensorflow/compiler/xla/hlo/evaluator/hlo_evaluator.h +14 -13
  186. tensorflow/include/tensorflow/compiler/xla/hlo/evaluator/hlo_evaluator_interface.h +64 -0
  187. tensorflow/include/tensorflow/compiler/xla/hlo/ir/dfs_hlo_visitor.h +1 -0
  188. tensorflow/include/tensorflow/compiler/xla/hlo/ir/dfs_hlo_visitor_with_default.h +3 -0
  189. tensorflow/include/tensorflow/compiler/xla/hlo/ir/hlo_instruction.h +11 -7
  190. tensorflow/include/tensorflow/compiler/xla/hlo/ir/hlo_instructions.h +70 -22
  191. tensorflow/include/tensorflow/compiler/xla/hlo/ir/hlo_module.h +70 -23
  192. tensorflow/include/tensorflow/compiler/xla/hlo/ir/hlo_opcode.h +2 -1
  193. tensorflow/include/tensorflow/compiler/xla/hlo/ir/hlo_original_value.h +1 -1
  194. tensorflow/include/tensorflow/compiler/xla/hlo/ir/hlo_sharding.h +24 -9
  195. tensorflow/include/tensorflow/compiler/xla/hlo/parser/hlo_lexer.h +0 -2
  196. tensorflow/include/tensorflow/compiler/xla/hlo/transforms/simplifiers/hlo_memory_scheduler.h +3 -14
  197. tensorflow/include/tensorflow/compiler/xla/hlo/transforms/simplifiers/reduce_window_resizer.h +43 -0
  198. tensorflow/include/tensorflow/compiler/xla/hlo/transforms/simplifiers/reduce_window_rewriter.h +0 -4
  199. tensorflow/include/tensorflow/compiler/xla/hlo/transforms/simplifiers/reduce_window_util.h +45 -0
  200. tensorflow/include/tensorflow/compiler/xla/hlo/translate/mhlo_to_hlo/attribute_exporter.h +33 -0
  201. tensorflow/include/tensorflow/compiler/xla/hlo/translate/register.h +31 -0
  202. tensorflow/include/tensorflow/compiler/xla/hlo/translate/stablehlo.h +0 -3
  203. tensorflow/include/tensorflow/compiler/xla/hlo/utils/hlo_sharding_util.h +5 -0
  204. tensorflow/include/tensorflow/compiler/xla/iterator_util.h +4 -0
  205. tensorflow/include/tensorflow/compiler/xla/pjrt/c/pjrt_c_api.h +17 -2
  206. tensorflow/include/tensorflow/compiler/xla/pjrt/c/pjrt_c_api_wrapper_impl.h +2 -0
  207. tensorflow/include/tensorflow/compiler/xla/pjrt/cpu/raw_buffer.h +3 -16
  208. tensorflow/include/tensorflow/compiler/xla/pjrt/device_event.h +17 -5
  209. tensorflow/include/tensorflow/compiler/xla/pjrt/distributed/client.h +4 -0
  210. tensorflow/include/tensorflow/compiler/xla/pjrt/distributed/protocol.pb.h +44 -44
  211. tensorflow/include/tensorflow/compiler/xla/pjrt/extensions/executable_metadata/executable_metadata_extension.h +83 -0
  212. tensorflow/include/tensorflow/compiler/xla/pjrt/gpu/gpu_helpers.h +3 -3
  213. tensorflow/include/tensorflow/compiler/xla/pjrt/gpu/gpu_topology.h +16 -12
  214. tensorflow/include/tensorflow/compiler/xla/pjrt/gpu/gpu_topology.pb.h +42 -42
  215. tensorflow/include/tensorflow/compiler/xla/pjrt/gpu/se_gpu_pjrt_client.h +18 -12
  216. tensorflow/include/tensorflow/compiler/xla/pjrt/gpu/se_gpu_topology_description.h +2 -1
  217. tensorflow/include/tensorflow/compiler/xla/pjrt/mlir_to_hlo.h +8 -13
  218. tensorflow/include/tensorflow/compiler/xla/pjrt/pjrt_c_api_client.h +13 -4
  219. tensorflow/include/tensorflow/compiler/xla/pjrt/pjrt_compiler.h +22 -0
  220. tensorflow/include/tensorflow/compiler/xla/pjrt/pjrt_executable.h +4 -0
  221. tensorflow/include/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.h +43 -28
  222. tensorflow/include/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_device_description.h +5 -2
  223. tensorflow/include/tensorflow/compiler/xla/pjrt/plugin/xla_gpu/xla_gpu_client_options.h +1 -1
  224. tensorflow/include/tensorflow/compiler/xla/pjrt/proto/compile_options.pb.h +37 -1
  225. tensorflow/include/tensorflow/compiler/xla/pjrt/proto/topology_description.pb.h +626 -0
  226. tensorflow/include/tensorflow/compiler/xla/pjrt/se_raw_buffer.h +117 -0
  227. tensorflow/include/tensorflow/compiler/xla/pjrt/thread_pool_async_work_runner.h +29 -0
  228. tensorflow/include/tensorflow/compiler/xla/pjrt/tracked_device_buffer.h +24 -22
  229. tensorflow/include/tensorflow/compiler/xla/service/buffer_assignment.h +15 -10
  230. tensorflow/include/tensorflow/compiler/xla/service/call_inliner.h +4 -6
  231. tensorflow/include/tensorflow/compiler/xla/service/compiler.h +22 -3
  232. tensorflow/include/tensorflow/compiler/xla/service/cpu/backend_config.pb.h +37 -1
  233. tensorflow/include/tensorflow/compiler/xla/service/cpu/cpu_compiler.h +0 -8
  234. tensorflow/include/tensorflow/compiler/xla/service/cpu/cpu_executable.h +7 -4
  235. tensorflow/include/tensorflow/compiler/xla/service/cpu/cpu_float_support.h +1 -1
  236. tensorflow/include/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.h +2 -2
  237. tensorflow/include/tensorflow/compiler/xla/service/cpu/cpu_multi_output_fusion.h +43 -0
  238. tensorflow/include/tensorflow/compiler/xla/service/cpu/cpu_options.h +6 -0
  239. tensorflow/include/tensorflow/compiler/xla/service/cpu/onednn_memory_util.h +1 -3
  240. tensorflow/include/tensorflow/compiler/xla/service/cpu/parallel_fusion_emitter.h +79 -0
  241. tensorflow/include/tensorflow/compiler/xla/service/cpu/thunk_emitter.h +7 -3
  242. tensorflow/include/tensorflow/compiler/xla/service/gpu/ir_emission_utils.h +0 -3
  243. tensorflow/include/tensorflow/compiler/xla/service/hlo.pb.h +2 -67
  244. tensorflow/include/tensorflow/compiler/xla/service/hlo_buffer.h +7 -2
  245. tensorflow/include/tensorflow/compiler/xla/service/hlo_cost_analysis.h +1 -0
  246. tensorflow/include/tensorflow/compiler/xla/service/hlo_creation_utils.h +1 -4
  247. tensorflow/include/tensorflow/compiler/xla/service/hlo_verifier.h +1 -4
  248. tensorflow/include/tensorflow/compiler/xla/service/instruction_fusion.h +2 -0
  249. tensorflow/include/tensorflow/compiler/xla/service/mapped_ptr_container_sorter.h +1 -1
  250. tensorflow/include/tensorflow/compiler/xla/service/matmul_indexing_utils.h +4 -1
  251. tensorflow/include/tensorflow/compiler/xla/service/multi_output_fusion.h +246 -0
  252. tensorflow/include/tensorflow/compiler/xla/service/shape_inference.h +1 -7
  253. tensorflow/include/tensorflow/compiler/xla/service/spmd/shardy/constants.h +1 -8
  254. tensorflow/include/{xla/service/spmd/shardy/round_trip_common/import_uninlineable_func_calls.h → tensorflow/compiler/xla/service/spmd/shardy/round_trip_common/import_func_calls.h} +16 -7
  255. tensorflow/include/tensorflow/compiler/xla/service/spmd/shardy/round_trip_common/pipeline_passes.h +2 -6
  256. tensorflow/include/tensorflow/compiler/xla/service/spmd/shardy/sdy_round_trip/export_shardy_attrs.h +2 -1
  257. tensorflow/include/tensorflow/compiler/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.h +2 -1
  258. tensorflow/include/tensorflow/compiler/xla/service/spmd/shardy/sdy_round_trip/pipelines.h +4 -9
  259. tensorflow/include/tensorflow/compiler/xla/service/spmd/shardy/sdy_round_trip/shard_map_export.h +5 -4
  260. tensorflow/include/tensorflow/compiler/xla/service/spmd/shardy/stablehlo_round_trip/shard_map_export.h +3 -1
  261. tensorflow/include/tensorflow/compiler/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_export.h +6 -0
  262. tensorflow/include/tensorflow/compiler/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import.h +12 -11
  263. tensorflow/include/tensorflow/compiler/xla/service/spmd/spmd_partitioner.h +10 -5
  264. tensorflow/include/tensorflow/compiler/xla/service/spmd/stateful_rng_spmd_partitioner.h +9 -3
  265. tensorflow/include/tensorflow/compiler/xla/stream_executor/command_buffer.h +8 -5
  266. tensorflow/include/tensorflow/compiler/xla/stream_executor/device_description.h +3 -1
  267. tensorflow/include/tensorflow/compiler/xla/stream_executor/gpu/tma_metadata.h +13 -0
  268. tensorflow/include/tensorflow/compiler/xla/stream_executor/stream.h +82 -21
  269. tensorflow/include/tensorflow/compiler/xla/stream_executor/stream_executor.h +61 -0
  270. tensorflow/include/tensorflow/compiler/xla/tsl/concurrency/async_value.h +43 -4
  271. tensorflow/include/tensorflow/compiler/xla/tsl/concurrency/async_value_ref.h +2 -9
  272. tensorflow/include/tensorflow/compiler/xla/tsl/distributed_runtime/coordination/coordination_service.h +0 -7
  273. tensorflow/include/tensorflow/compiler/xla/tsl/platform/status.h +0 -5
  274. tensorflow/include/tensorflow/compiler/xla/tsl/util/onednn_threadpool.h +7 -0
  275. tensorflow/include/tensorflow/compiler/xla/xla.pb.h +443 -223
  276. tensorflow/include/tensorflow/compiler/xla/xla_data.pb.h +32 -444
  277. tensorflow/include/tensorflow/core/common_runtime/eager/eager_executor.h +2 -1
  278. tensorflow/include/tensorflow/core/distributed_runtime/worker_interface.h +4 -3
  279. tensorflow/include/tensorflow/core/framework/node_properties.h +8 -8
  280. tensorflow/include/tensorflow/core/framework/op_kernel.h +41 -36
  281. tensorflow/include/tensorflow/core/ir/ops.h.inc +1 -1
  282. tensorflow/include/tensorflow/core/kernels/batching_util/batch_scheduler.h +4 -4
  283. tensorflow/include/tensorflow/core/kernels/batching_util/periodic_function.h +2 -2
  284. tensorflow/include/tensorflow/core/kernels/batching_util/shared_batch_scheduler.h +3 -2
  285. tensorflow/include/tensorflow/core/kernels/debug_ops.h +2 -2
  286. tensorflow/include/tensorflow/core/kernels/record_yielder.h +2 -2
  287. tensorflow/include/tensorflow/core/lib/jpeg/jpeg_mem.h +1 -1
  288. tensorflow/include/tensorflow/core/platform/status.h +0 -6
  289. tensorflow/include/tensorflow/core/protobuf/config.pb.h +153 -71
  290. tensorflow/include/tensorflow/core/protobuf/for_core_protos-descriptor-set.proto.bin +0 -0
  291. tensorflow/include/tensorflow/core/public/version.h +1 -1
  292. tensorflow/include/xla/backends/cpu/codegen/emitters/cpu_scatter_emitter.h +4 -3
  293. tensorflow/include/xla/backends/cpu/codegen/execution_engine.h +1 -1
  294. tensorflow/include/xla/backends/cpu/codegen/fusion_compiler.h +11 -5
  295. tensorflow/include/xla/backends/cpu/codegen/ir_compiler.h +1 -0
  296. tensorflow/include/xla/backends/cpu/{onednn_fusion.h → onednn_support.h} +4 -4
  297. tensorflow/include/xla/backends/cpu/runtime/thunk.h +108 -58
  298. tensorflow/include/xla/backends/cpu/runtime/thunk_executor.h +11 -20
  299. tensorflow/include/xla/backends/cpu/runtime/work_queue.h +3 -13
  300. tensorflow/include/xla/backends/cpu/runtime/xnnpack/xnn_convolution_thunk.h +1 -1
  301. tensorflow/include/xla/backends/cpu/runtime/xnnpack/xnn_dot_thunk.h +1 -1
  302. tensorflow/include/xla/backends/cpu/runtime/xnnpack/xnn_fusion_thunk.h +20 -19
  303. tensorflow/include/xla/backends/cpu/runtime/xnnpack/xnn_interop.h +45 -0
  304. tensorflow/include/xla/backends/cpu/runtime/xnnpack/xnn_threadpool.h +11 -27
  305. tensorflow/include/xla/backends/cpu/transforms/onednn_matcher.h +1 -1
  306. tensorflow/include/xla/backends/cpu/transforms/xnn_matcher.h +5 -5
  307. tensorflow/include/xla/backends/cpu/xnn_emitter.h +2 -2
  308. tensorflow/include/xla/backends/cpu/xnn_gemm_config.h +1 -0
  309. tensorflow/include/xla/backends/cpu/{xnn_fusion.h → xnn_support.h} +7 -8
  310. tensorflow/include/xla/backends/gpu/codegen/emitters/ir/_virtual_includes/xla_gpu_ops_inc_gen/xla_gpu_ops.h.inc +2 -2
  311. tensorflow/include/xla/backends/gpu/codegen/emitters/ir/xla_gpu_ops.h.inc +2 -2
  312. tensorflow/include/xla/client/executable_build_options.h +5 -0
  313. tensorflow/include/xla/codegen/emitters/concatenate_kernel_emitter.h +0 -1
  314. tensorflow/include/xla/codegen/emitters/ir/_virtual_includes/xla_ops_inc_gen/xla_ops.h.inc +1 -1
  315. tensorflow/include/xla/codegen/emitters/ir/xla_ops.h.inc +1 -1
  316. tensorflow/include/xla/codegen/emitters/loop_kernel_emitter.h +0 -1
  317. tensorflow/include/xla/codegen/emitters/transforms/pass_pipelines.h +28 -0
  318. tensorflow/include/xla/codegen/emitters/transforms/passes.h +1 -1
  319. tensorflow/include/xla/codegen/emitters/transforms/passes.h.inc +92 -92
  320. tensorflow/include/xla/codegen/emitters/utils.h +3 -1
  321. tensorflow/include/{tensorflow/compiler/xla/codegen/math → xla/codegen/intrinsic}/erf.h +4 -4
  322. tensorflow/include/{tensorflow/compiler/xla/codegen/math → xla/codegen/intrinsic}/exp.h +4 -4
  323. tensorflow/include/{tensorflow/compiler/xla/codegen/math → xla/codegen/intrinsic}/fptrunc.h +4 -4
  324. tensorflow/include/xla/codegen/{math → intrinsic}/intrinsic.h +29 -3
  325. tensorflow/include/xla/codegen/{math/math_compiler_lib.h → intrinsic/intrinsic_compiler_lib.h} +5 -5
  326. tensorflow/include/xla/codegen/{math → intrinsic}/ldexp.h +4 -4
  327. tensorflow/include/xla/codegen/{math → intrinsic}/log1p.h +4 -4
  328. tensorflow/include/{tensorflow/compiler/xla/codegen/math → xla/codegen/intrinsic}/rsqrt.h +9 -9
  329. tensorflow/include/{tensorflow/compiler/xla/codegen/math → xla/codegen/intrinsic}/string_interner.h +5 -5
  330. tensorflow/include/{tensorflow/compiler/xla/codegen/math → xla/codegen/intrinsic}/tanh.h +6 -9
  331. tensorflow/include/{tensorflow/compiler/xla/codegen/math → xla/codegen/intrinsic}/vec_name_mangler.h +5 -5
  332. tensorflow/include/{tensorflow/compiler/xla/codegen/math_lib.h → xla/codegen/intrinsic_lib.h} +16 -18
  333. tensorflow/include/xla/codegen/trace_pass_instrumentation.h +47 -0
  334. tensorflow/include/xla/core/collectives/communicator.h +10 -0
  335. tensorflow/include/xla/fp_util.h +3 -1
  336. tensorflow/include/xla/hlo/analysis/indexing_analysis.h +14 -2
  337. tensorflow/include/xla/hlo/builder/xla_builder.h +14 -0
  338. tensorflow/include/xla/hlo/evaluator/hlo_evaluator.h +14 -13
  339. tensorflow/include/xla/hlo/evaluator/hlo_evaluator_interface.h +64 -0
  340. tensorflow/include/xla/hlo/ir/dfs_hlo_visitor.h +1 -0
  341. tensorflow/include/xla/hlo/ir/dfs_hlo_visitor_with_default.h +3 -0
  342. tensorflow/include/xla/hlo/ir/hlo_instruction.h +11 -7
  343. tensorflow/include/xla/hlo/ir/hlo_instructions.h +70 -22
  344. tensorflow/include/xla/hlo/ir/hlo_module.h +70 -23
  345. tensorflow/include/xla/hlo/ir/hlo_opcode.h +2 -1
  346. tensorflow/include/xla/hlo/ir/hlo_original_value.h +1 -1
  347. tensorflow/include/xla/hlo/ir/hlo_sharding.h +24 -9
  348. tensorflow/include/xla/hlo/parser/hlo_lexer.h +0 -2
  349. tensorflow/include/xla/hlo/transforms/simplifiers/hlo_memory_scheduler.h +3 -14
  350. tensorflow/include/xla/hlo/transforms/simplifiers/reduce_window_resizer.h +43 -0
  351. tensorflow/include/xla/hlo/transforms/simplifiers/reduce_window_rewriter.h +0 -4
  352. tensorflow/include/xla/hlo/transforms/simplifiers/reduce_window_util.h +45 -0
  353. tensorflow/include/xla/hlo/translate/mhlo_to_hlo/attribute_exporter.h +33 -0
  354. tensorflow/include/xla/hlo/translate/register.h +31 -0
  355. tensorflow/include/xla/hlo/translate/stablehlo.h +0 -3
  356. tensorflow/include/xla/hlo/utils/hlo_sharding_util.h +5 -0
  357. tensorflow/include/xla/iterator_util.h +4 -0
  358. tensorflow/include/xla/pjrt/c/pjrt_c_api.h +17 -2
  359. tensorflow/include/xla/pjrt/c/pjrt_c_api_wrapper_impl.h +2 -0
  360. tensorflow/include/xla/pjrt/cpu/raw_buffer.h +3 -16
  361. tensorflow/include/xla/pjrt/device_event.h +17 -5
  362. tensorflow/include/xla/pjrt/distributed/client.h +4 -0
  363. tensorflow/include/xla/pjrt/distributed/protocol.pb.h +44 -44
  364. tensorflow/include/xla/pjrt/extensions/executable_metadata/executable_metadata_extension.h +83 -0
  365. tensorflow/include/xla/pjrt/gpu/gpu_helpers.h +3 -3
  366. tensorflow/include/xla/pjrt/gpu/gpu_topology.h +16 -12
  367. tensorflow/include/xla/pjrt/gpu/gpu_topology.pb.h +42 -42
  368. tensorflow/include/xla/pjrt/gpu/se_gpu_pjrt_client.h +18 -12
  369. tensorflow/include/xla/pjrt/gpu/se_gpu_topology_description.h +2 -1
  370. tensorflow/include/xla/pjrt/mlir_to_hlo.h +8 -13
  371. tensorflow/include/xla/pjrt/pjrt_c_api_client.h +13 -4
  372. tensorflow/include/xla/pjrt/pjrt_compiler.h +22 -0
  373. tensorflow/include/xla/pjrt/pjrt_executable.h +4 -0
  374. tensorflow/include/xla/pjrt/pjrt_stream_executor_client.h +43 -28
  375. tensorflow/include/xla/pjrt/pjrt_stream_executor_device_description.h +5 -2
  376. tensorflow/include/xla/pjrt/plugin/xla_gpu/xla_gpu_client_options.h +1 -1
  377. tensorflow/include/xla/pjrt/proto/compile_options.pb.h +37 -1
  378. tensorflow/include/xla/pjrt/proto/topology_description.pb.h +626 -0
  379. tensorflow/include/xla/pjrt/se_raw_buffer.h +117 -0
  380. tensorflow/include/xla/pjrt/thread_pool_async_work_runner.h +29 -0
  381. tensorflow/include/xla/pjrt/tracked_device_buffer.h +24 -22
  382. tensorflow/include/xla/service/buffer_assignment.h +15 -10
  383. tensorflow/include/xla/service/call_inliner.h +4 -6
  384. tensorflow/include/xla/service/compiler.h +22 -3
  385. tensorflow/include/xla/service/cpu/backend_config.pb.h +37 -1
  386. tensorflow/include/xla/service/cpu/cpu_compiler.h +0 -8
  387. tensorflow/include/xla/service/cpu/cpu_executable.h +7 -4
  388. tensorflow/include/xla/service/cpu/cpu_float_support.h +1 -1
  389. tensorflow/include/xla/service/cpu/cpu_instruction_fusion.h +2 -2
  390. tensorflow/include/xla/service/cpu/cpu_multi_output_fusion.h +43 -0
  391. tensorflow/include/xla/service/cpu/cpu_options.h +6 -0
  392. tensorflow/include/xla/service/cpu/onednn_memory_util.h +1 -3
  393. tensorflow/include/xla/service/cpu/parallel_fusion_emitter.h +79 -0
  394. tensorflow/include/xla/service/cpu/thunk_emitter.h +7 -3
  395. tensorflow/include/xla/service/gpu/ir_emission_utils.h +0 -3
  396. tensorflow/include/xla/service/hlo.pb.h +2 -67
  397. tensorflow/include/xla/service/hlo_buffer.h +7 -2
  398. tensorflow/include/xla/service/hlo_cost_analysis.h +1 -0
  399. tensorflow/include/xla/service/hlo_creation_utils.h +1 -4
  400. tensorflow/include/xla/service/hlo_verifier.h +1 -4
  401. tensorflow/include/xla/service/instruction_fusion.h +2 -0
  402. tensorflow/include/xla/service/mapped_ptr_container_sorter.h +1 -1
  403. tensorflow/include/xla/service/matmul_indexing_utils.h +4 -1
  404. tensorflow/include/xla/service/multi_output_fusion.h +246 -0
  405. tensorflow/include/xla/service/shape_inference.h +1 -7
  406. tensorflow/include/xla/service/spmd/shardy/constants.h +1 -8
  407. tensorflow/include/{tensorflow/compiler/xla/service/spmd/shardy/round_trip_common/import_uninlineable_func_calls.h → xla/service/spmd/shardy/round_trip_common/import_func_calls.h} +16 -7
  408. tensorflow/include/xla/service/spmd/shardy/round_trip_common/pipeline_passes.h +2 -6
  409. tensorflow/include/xla/service/spmd/shardy/sdy_round_trip/export_shardy_attrs.h +2 -1
  410. tensorflow/include/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.h +2 -1
  411. tensorflow/include/xla/service/spmd/shardy/sdy_round_trip/pipelines.h +4 -9
  412. tensorflow/include/xla/service/spmd/shardy/sdy_round_trip/shard_map_export.h +5 -4
  413. tensorflow/include/xla/service/spmd/shardy/stablehlo_round_trip/shard_map_export.h +3 -1
  414. tensorflow/include/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_export.h +6 -0
  415. tensorflow/include/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import.h +12 -11
  416. tensorflow/include/xla/service/spmd/spmd_partitioner.h +10 -5
  417. tensorflow/include/xla/service/spmd/stateful_rng_spmd_partitioner.h +9 -3
  418. tensorflow/include/xla/stream_executor/command_buffer.h +8 -5
  419. tensorflow/include/xla/stream_executor/device_description.h +3 -1
  420. tensorflow/include/xla/stream_executor/gpu/tma_metadata.h +13 -0
  421. tensorflow/include/xla/stream_executor/stream.h +82 -21
  422. tensorflow/include/xla/stream_executor/stream_executor.h +61 -0
  423. tensorflow/include/xla/tsl/concurrency/async_value.h +43 -4
  424. tensorflow/include/xla/tsl/concurrency/async_value_ref.h +2 -9
  425. tensorflow/include/xla/tsl/distributed_runtime/coordination/coordination_service.h +0 -7
  426. tensorflow/include/xla/tsl/platform/status.h +0 -5
  427. tensorflow/include/xla/tsl/util/onednn_threadpool.h +7 -0
  428. tensorflow/include/xla/xla.pb.h +443 -223
  429. tensorflow/include/xla/xla_data.pb.h +32 -444
  430. tensorflow/libtensorflow_cc.2.dylib +0 -0
  431. tensorflow/libtensorflow_framework.2.dylib +0 -0
  432. tensorflow/lite/experimental/microfrontend/python/ops/_audio_microfrontend_op.so +0 -0
  433. tensorflow/lite/python/libpywrap_tflite_common.dylib +0 -0
  434. tensorflow/python/_pywrap_tfcompile.so +0 -0
  435. tensorflow/python/_pywrap_tfe.so +0 -0
  436. tensorflow/python/client/_pywrap_tf_session.so +0 -0
  437. tensorflow/python/compat/compat.py +1 -1
  438. tensorflow/python/distribute/multi_process_runner.py +95 -9
  439. tensorflow/python/keras/regularizers.py +6 -4
  440. tensorflow/python/lib_pywrap_tensorflow_common.dylib +0 -0
  441. tensorflow/python/ops/gen_list_ops.py +8 -6
  442. tensorflow/python/ops/string_ops.py +6 -6
  443. tensorflow/python/profiler/internal/_pywrap_profiler.so +0 -0
  444. tensorflow/python/profiler/internal/_pywrap_profiler_plugin.so +0 -0
  445. tensorflow/python/tpu/_pywrap_tpu_embedding.so +0 -0
  446. tensorflow/tools/pip_package/setup.py +2 -2
  447. tensorflow/xla_aot_runtime_src/tensorflow/compiler/tf2xla/xla_compiled_cpu_function.cc +1 -9
  448. tensorflow/xla_aot_runtime_src/xla/tsl/concurrency/async_value.cc +1 -36
  449. tensorflow/include/tensorflow/compiler/xla/backends/cpu/runtime/parallel_loop_runner.h +0 -490
  450. tensorflow/include/tensorflow/compiler/xla/pjrt/gpu/raw_buffer.h +0 -63
  451. tensorflow/include/xla/backends/cpu/runtime/parallel_loop_runner.h +0 -490
  452. tensorflow/include/xla/pjrt/gpu/raw_buffer.h +0 -63
  453. {ai_edge_tensorflow-2.21.0.dev20250805.dist-info → ai_edge_tensorflow-2.21.0.dev20250818.dist-info}/WHEEL +0 -0
  454. {ai_edge_tensorflow-2.21.0.dev20250805.dist-info → ai_edge_tensorflow-2.21.0.dev20250818.dist-info}/entry_points.txt +0 -0
  455. {ai_edge_tensorflow-2.21.0.dev20250805.dist-info → ai_edge_tensorflow-2.21.0.dev20250818.dist-info}/top_level.txt +0 -0
{ai_edge_tensorflow-2.21.0.dev20250805.dist-info → ai_edge_tensorflow-2.21.0.dev20250818.dist-info}/METADATA RENAMED
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: ai_edge_tensorflow
3
- Version: 2.21.0.dev20250805
3
+ Version: 2.21.0.dev20250818
4
4
  Summary: TensorFlow is an open source machine learning framework for everyone.
5
5
  Home-page: https://www.tensorflow.org/
6
6
  Download-URL: https://github.com/tensorflow/tensorflow/tags
@@ -62,7 +62,7 @@ Requires-Dist: nvidia-cufft-cu12<12.0,>=11.2.3.61; extra == "and-cuda"
62
62
  Requires-Dist: nvidia-curand-cu12<11.0,>=10.3.6.82; extra == "and-cuda"
63
63
  Requires-Dist: nvidia-cusolver-cu12<12.0,>=11.6.3.83; extra == "and-cuda"
64
64
  Requires-Dist: nvidia-cusparse-cu12<13.0,>=12.5.1.3; extra == "and-cuda"
65
- Requires-Dist: nvidia-nccl-cu12<3.0,>=2.26.5; extra == "and-cuda"
65
+ Requires-Dist: nvidia-nccl-cu12<3.0,>=2.27.7; extra == "and-cuda"
66
66
  Requires-Dist: nvidia-nvjitlink-cu12<13.0,>=12.5.82; extra == "and-cuda"
67
67
  Provides-Extra: gcs-filesystem
68
68
  Requires-Dist: tensorflow-io-gcs-filesystem>=0.23.1; (sys_platform != "win32" and python_version < "3.13") and extra == "gcs-filesystem"