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