tf-nightly-cpu 2.20.0.dev20250220__cp312-cp312-win_amd64.whl → 2.20.0.dev20250222__cp312-cp312-win_amd64.whl

Sign up to get free protection for your applications and to get access to all the features.
Files changed (130) hide show
  1. tensorflow/_api/v2/compat/v1/summary/__init__.py +2 -2
  2. tensorflow/_api/v2/compat/v1/tpu/experimental/embedding/__init__.py +2 -2
  3. tensorflow/_api/v2/compat/v2/summary/__init__.py +10 -10
  4. tensorflow/_api/v2/compat/v2/summary/experimental/__init__.py +4 -4
  5. tensorflow/_api/v2/compat/v2/tpu/experimental/embedding/__init__.py +2 -2
  6. tensorflow/_api/v2/summary/__init__.py +10 -10
  7. tensorflow/_api/v2/summary/experimental/__init__.py +4 -4
  8. tensorflow/_api/v2/tpu/experimental/embedding/__init__.py +2 -2
  9. tensorflow/compiler/mlir/stablehlo/stablehlo_extension.pyd +0 -0
  10. tensorflow/compiler/tf2tensorrt/_pywrap_py_utils.pyd +0 -0
  11. tensorflow/compiler/tf2xla/ops/_xla_ops.so +0 -0
  12. tensorflow/include/external/llvm-project/mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h +12 -0
  13. tensorflow/include/external/llvm-project/mlir/include/mlir/Dialect/Math/IR/MathOps.h.inc +4 -0
  14. tensorflow/include/external/shardy/shardy/dialect/sdy/transforms/propagation/aggressive_factor_propagation.h +9 -0
  15. tensorflow/include/external/stablehlo/_virtual_includes/stablehlo_pass_utils/stablehlo/transforms/PassUtils.h +7 -0
  16. tensorflow/include/external/stablehlo/_virtual_includes/stablehlo_passes/stablehlo/transforms/PassUtils.h +7 -0
  17. tensorflow/include/external/stablehlo/_virtual_includes/version/stablehlo/dialect/Version.h +1 -1
  18. tensorflow/include/external/stablehlo/stablehlo/dialect/Version.h +1 -1
  19. tensorflow/include/external/stablehlo/stablehlo/transforms/PassUtils.h +7 -0
  20. tensorflow/include/tensorflow/compiler/xla/backends/cpu/codegen/kernel_api_ir_builder.h +3 -2
  21. tensorflow/include/tensorflow/compiler/xla/backends/cpu/runtime/convolution_thunk_internal.h +8 -10
  22. tensorflow/include/tensorflow/compiler/xla/backends/cpu/runtime/kernel_thunk.h +9 -3
  23. tensorflow/include/tensorflow/compiler/xla/backends/cpu/runtime/work_queue.h +81 -19
  24. tensorflow/include/tensorflow/compiler/xla/codegen/kernel_spec.h +24 -7
  25. tensorflow/include/tensorflow/compiler/xla/hlo/ir/hlo_casting_utils.h +0 -44
  26. tensorflow/include/tensorflow/compiler/xla/hlo/ir/hlo_instruction.h +12 -0
  27. tensorflow/include/tensorflow/compiler/xla/mlir_hlo/_virtual_includes/stablehlo_extension_pass_inc_gen/stablehlo_ext/transforms/passes.h.inc +149 -4
  28. tensorflow/include/tensorflow/compiler/xla/mlir_hlo/stablehlo_ext/transforms/passes.h.inc +149 -4
  29. tensorflow/include/tensorflow/compiler/xla/pjrt/distributed/client.h +5 -0
  30. tensorflow/include/tensorflow/compiler/xla/pjrt/gpu/se_gpu_pjrt_client.h +1 -92
  31. tensorflow/include/tensorflow/compiler/xla/pjrt/gpu/se_gpu_topology_description.h +126 -0
  32. tensorflow/include/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.h +1 -49
  33. tensorflow/include/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_device_description.h +75 -0
  34. tensorflow/include/tensorflow/compiler/xla/pjrt/plugin/xla_cpu/cpu_execute_options.h +57 -0
  35. tensorflow/include/tensorflow/compiler/xla/pjrt/plugin/xla_cpu/cpu_topology.h +4 -0
  36. tensorflow/include/tensorflow/compiler/xla/service/constant_value.h +1 -0
  37. tensorflow/include/tensorflow/compiler/xla/service/hlo_module_util.h +52 -1
  38. tensorflow/include/tensorflow/compiler/xla/service/hlo_proto_util.h +0 -12
  39. tensorflow/include/tensorflow/compiler/xla/tsl/concurrency/async_value.h +50 -21
  40. tensorflow/include/tensorflow/compiler/xla/tsl/framework/convolution/eigen_spatial_convolutions-inl.h +5 -5
  41. tensorflow/include/tensorflow/core/kernels/data/experimental/random_access_ops.h +0 -2
  42. tensorflow/include/tensorflow/core/kernels/eigen_attention.h +4 -4
  43. tensorflow/include/tensorflow/core/kernels/eigen_backward_cuboid_convolutions.h +6 -6
  44. tensorflow/include/tensorflow/core/kernels/eigen_backward_spatial_convolutions.h +10 -8
  45. tensorflow/include/tensorflow/core/kernels/eigen_cuboid_convolution.h +6 -6
  46. tensorflow/include/tensorflow/core/kernels/eigen_pooling.h +12 -12
  47. tensorflow/include/tensorflow/core/public/release_version.h +39 -0
  48. tensorflow/include/tensorflow/core/public/version.h +112 -127
  49. tensorflow/include/tensorflow/python/eager/pywrap_tfe.h +1 -1
  50. tensorflow/include/xla/backends/cpu/codegen/kernel_api_ir_builder.h +3 -2
  51. tensorflow/include/xla/backends/cpu/runtime/convolution_thunk_internal.h +8 -10
  52. tensorflow/include/xla/backends/cpu/runtime/kernel_thunk.h +9 -3
  53. tensorflow/include/xla/backends/cpu/runtime/work_queue.h +81 -19
  54. tensorflow/include/xla/codegen/kernel_spec.h +24 -7
  55. tensorflow/include/xla/hlo/ir/hlo_casting_utils.h +0 -44
  56. tensorflow/include/xla/hlo/ir/hlo_instruction.h +12 -0
  57. tensorflow/include/xla/mlir_hlo/_virtual_includes/stablehlo_extension_pass_inc_gen/stablehlo_ext/transforms/passes.h.inc +149 -4
  58. tensorflow/include/xla/mlir_hlo/stablehlo_ext/transforms/passes.h.inc +149 -4
  59. tensorflow/include/xla/pjrt/distributed/client.h +5 -0
  60. tensorflow/include/xla/pjrt/gpu/se_gpu_pjrt_client.h +1 -92
  61. tensorflow/include/xla/pjrt/gpu/se_gpu_topology_description.h +126 -0
  62. tensorflow/include/xla/pjrt/pjrt_stream_executor_client.h +1 -49
  63. tensorflow/include/xla/pjrt/pjrt_stream_executor_device_description.h +75 -0
  64. tensorflow/include/xla/pjrt/plugin/xla_cpu/cpu_execute_options.h +57 -0
  65. tensorflow/include/xla/pjrt/plugin/xla_cpu/cpu_topology.h +4 -0
  66. tensorflow/include/xla/service/constant_value.h +1 -0
  67. tensorflow/include/xla/service/hlo_module_util.h +52 -1
  68. tensorflow/include/xla/service/hlo_proto_util.h +0 -12
  69. tensorflow/include/xla/tsl/concurrency/async_value.h +50 -21
  70. tensorflow/include/xla/tsl/framework/convolution/eigen_spatial_convolutions-inl.h +5 -5
  71. tensorflow/lite/experimental/microfrontend/python/ops/_audio_microfrontend_op.so +0 -0
  72. tensorflow/lite/python/analyzer_wrapper/_pywrap_analyzer_wrapper.pyd +0 -0
  73. tensorflow/lite/python/interpreter_wrapper/_pywrap_tensorflow_interpreter_wrapper.pyd +0 -0
  74. tensorflow/lite/python/optimize/_pywrap_tensorflow_lite_calibration_wrapper.pyd +0 -0
  75. tensorflow/python/_pywrap_dtensor_device.pyd +0 -0
  76. tensorflow/python/_pywrap_mlir.pyd +0 -0
  77. tensorflow/python/_pywrap_parallel_device.pyd +0 -0
  78. tensorflow/python/_pywrap_quantize_training.pyd +0 -0
  79. tensorflow/python/_pywrap_tensorflow_internal.pyd +0 -0
  80. tensorflow/python/_pywrap_tfcompile.pyd +0 -0
  81. tensorflow/python/_pywrap_tfe.pyd +0 -0
  82. tensorflow/python/client/_pywrap_debug_events_writer.pyd +0 -0
  83. tensorflow/python/client/_pywrap_device_lib.pyd +0 -0
  84. tensorflow/python/client/_pywrap_events_writer.pyd +0 -0
  85. tensorflow/python/client/_pywrap_tf_session.pyd +0 -0
  86. tensorflow/python/compat/compat.py +1 -1
  87. tensorflow/python/data/experimental/service/_pywrap_server_lib.pyd +0 -0
  88. tensorflow/python/data/experimental/service/_pywrap_utils_exp.pyd +0 -0
  89. tensorflow/python/eager/imperative_grad.py +5 -5
  90. tensorflow/python/eager/polymorphic_function/atomic_function.py +1 -1
  91. tensorflow/python/eager/polymorphic_function/compiler_ir.py +1 -1
  92. tensorflow/python/eager/polymorphic_function/polymorphic_function.py +45 -41
  93. tensorflow/python/eager/tape.py +2 -2
  94. tensorflow/python/framework/_dtypes.pyd +0 -0
  95. tensorflow/python/framework/_op_def_library_pybind.pyd +0 -0
  96. tensorflow/python/framework/_op_def_registry.pyd +0 -0
  97. tensorflow/python/framework/_proto_comparators.pyd +0 -0
  98. tensorflow/python/framework/_pywrap_python_op_gen.pyd +0 -0
  99. tensorflow/python/framework/_test_metrics_util.pyd +0 -0
  100. tensorflow/python/grappler/_pywrap_tf_cluster.pyd +0 -0
  101. tensorflow/python/grappler/_pywrap_tf_item.pyd +0 -0
  102. tensorflow/python/grappler/_pywrap_tf_optimizer.pyd +0 -0
  103. tensorflow/python/lib/core/_pywrap_py_func.pyd +0 -0
  104. tensorflow/python/lib/io/_pywrap_file_io.pyd +0 -0
  105. tensorflow/python/lib/io/_pywrap_record_io.pyd +0 -0
  106. tensorflow/python/ops/summary_ops_v2.py +5 -1
  107. tensorflow/python/platform/_pywrap_tf2.pyd +0 -0
  108. tensorflow/python/profiler/internal/_pywrap_profiler.pyd +0 -0
  109. tensorflow/python/profiler/internal/_pywrap_profiler_plugin.pyd +0 -0
  110. tensorflow/python/saved_model/pywrap_saved_model.pyd +0 -0
  111. tensorflow/python/tpu/_pywrap_sparse_core_layout.pyd +0 -0
  112. tensorflow/python/tpu/_pywrap_tpu_embedding.pyd +0 -0
  113. tensorflow/python/tpu/tpu_embedding_v3.py +14 -7
  114. tensorflow/python/tpu/tpu_embedding_v3_checkpoint_adapter.py +10 -1
  115. tensorflow/python/util/_pywrap_checkpoint_reader.pyd +0 -0
  116. tensorflow/python/util/_pywrap_kernel_registry.pyd +0 -0
  117. tensorflow/python/util/_pywrap_stat_summarizer.pyd +0 -0
  118. tensorflow/python/util/_pywrap_tfprof.pyd +0 -0
  119. tensorflow/python/util/_pywrap_transform_graph.pyd +0 -0
  120. tensorflow/python/util/_pywrap_utils.pyd +0 -0
  121. tensorflow/python/util/_tf_stack.pyd +0 -0
  122. tensorflow/tools/pip_package/setup.py +2 -2
  123. tensorflow/xla_aot_runtime_src/xla/tsl/concurrency/async_value.cc +26 -51
  124. {tf_nightly_cpu-2.20.0.dev20250220.dist-info → tf_nightly_cpu-2.20.0.dev20250222.dist-info}/METADATA +1 -1
  125. {tf_nightly_cpu-2.20.0.dev20250220.dist-info → tf_nightly_cpu-2.20.0.dev20250222.dist-info}/RECORD +128 -123
  126. tensorflow/include/tensorflow/compiler/xla/backends/cpu/runtime/concurrency.h +0 -77
  127. tensorflow/include/xla/backends/cpu/runtime/concurrency.h +0 -77
  128. {tf_nightly_cpu-2.20.0.dev20250220.dist-info → tf_nightly_cpu-2.20.0.dev20250222.dist-info}/WHEEL +0 -0
  129. {tf_nightly_cpu-2.20.0.dev20250220.dist-info → tf_nightly_cpu-2.20.0.dev20250222.dist-info}/entry_points.txt +0 -0
  130. {tf_nightly_cpu-2.20.0.dev20250220.dist-info → tf_nightly_cpu-2.20.0.dev20250222.dist-info}/top_level.txt +0 -0
@@ -1,77 +0,0 @@
1
- /* Copyright 2024 The OpenXLA Authors.
2
-
3
- Licensed under the Apache License, Version 2.0 (the "License");
4
- you may not use this file except in compliance with the License.
5
- You may obtain a copy of the License at
6
-
7
- http://www.apache.org/licenses/LICENSE-2.0
8
-
9
- Unless required by applicable law or agreed to in writing, software
10
- distributed under the License is distributed on an "AS IS" BASIS,
11
- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
- See the License for the specific language governing permissions and
13
- limitations under the License.
14
- ==============================================================================*/
15
-
16
- #ifndef XLA_BACKENDS_CPU_RUNTIME_CONCURRENCY_H_
17
- #define XLA_BACKENDS_CPU_RUNTIME_CONCURRENCY_H_
18
-
19
- #include <cstdint>
20
- #include <functional>
21
- #include <memory>
22
- #include <type_traits>
23
-
24
- #include "tsl/platform/logging.h"
25
-
26
- #define EIGEN_USE_THREADS
27
- #include "unsupported/Eigen/CXX11/Tensor"
28
- #include "unsupported/Eigen/CXX11/ThreadPool"
29
-
30
- namespace xla::cpu {
31
-
32
- // A collection of concurrency primitives for use in the XLA CPU runtime.
33
-
34
- // Schedules `n` tasks on the `intra_op_threadpool`, calling `F` for each index
35
- // in the [0, n) range. Returns immediately after scheduling all tasks. It's a
36
- // caller's responsibility to wait for all tasks to finish.
37
- template <typename F,
38
- std::enable_if_t<std::is_invocable_v<F, int64_t>>* = nullptr>
39
- void ScheduleAll(const Eigen::ThreadPoolDevice* intra_op_threadpool, int64_t n,
40
- F&& f) {
41
- DCHECK(n >= 0) << "n must be non-negative";
42
-
43
- // Short-circuit the case of no tasks.
44
- if (n == 0) return;
45
-
46
- // Short-circuit the case of a single task.
47
- if (n == 1) {
48
- f(0);
49
- return;
50
- }
51
-
52
- // Heap-allocated state that manages concurrent execution of `f`.
53
- struct State {
54
- State(const Eigen::ThreadPoolDevice* intra_op_threadpool, F&& f)
55
- : intra_op_threadpool(intra_op_threadpool), f(std::forward<F>(f)) {}
56
-
57
- void Execute(std::shared_ptr<State> self, int64_t start, int64_t end) {
58
- while (end - start > 1) {
59
- uint64_t mid = (start + end) / 2;
60
- intra_op_threadpool->getPool()->Schedule(
61
- std::bind(&State::Execute, this, self, mid, end));
62
- end = mid;
63
- }
64
- f(start);
65
- }
66
-
67
- const Eigen::ThreadPoolDevice* intra_op_threadpool;
68
- F f;
69
- };
70
-
71
- auto s = std::make_shared<State>(intra_op_threadpool, std::forward<F>(f));
72
- s->Execute(std::move(s), 0, n);
73
- }
74
-
75
- } // namespace xla::cpu
76
-
77
- #endif // XLA_BACKENDS_CPU_RUNTIME_CONCURRENCY_H_
@@ -1,77 +0,0 @@
1
- /* Copyright 2024 The OpenXLA Authors.
2
-
3
- Licensed under the Apache License, Version 2.0 (the "License");
4
- you may not use this file except in compliance with the License.
5
- You may obtain a copy of the License at
6
-
7
- http://www.apache.org/licenses/LICENSE-2.0
8
-
9
- Unless required by applicable law or agreed to in writing, software
10
- distributed under the License is distributed on an "AS IS" BASIS,
11
- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
- See the License for the specific language governing permissions and
13
- limitations under the License.
14
- ==============================================================================*/
15
-
16
- #ifndef XLA_BACKENDS_CPU_RUNTIME_CONCURRENCY_H_
17
- #define XLA_BACKENDS_CPU_RUNTIME_CONCURRENCY_H_
18
-
19
- #include <cstdint>
20
- #include <functional>
21
- #include <memory>
22
- #include <type_traits>
23
-
24
- #include "tsl/platform/logging.h"
25
-
26
- #define EIGEN_USE_THREADS
27
- #include "unsupported/Eigen/CXX11/Tensor"
28
- #include "unsupported/Eigen/CXX11/ThreadPool"
29
-
30
- namespace xla::cpu {
31
-
32
- // A collection of concurrency primitives for use in the XLA CPU runtime.
33
-
34
- // Schedules `n` tasks on the `intra_op_threadpool`, calling `F` for each index
35
- // in the [0, n) range. Returns immediately after scheduling all tasks. It's a
36
- // caller's responsibility to wait for all tasks to finish.
37
- template <typename F,
38
- std::enable_if_t<std::is_invocable_v<F, int64_t>>* = nullptr>
39
- void ScheduleAll(const Eigen::ThreadPoolDevice* intra_op_threadpool, int64_t n,
40
- F&& f) {
41
- DCHECK(n >= 0) << "n must be non-negative";
42
-
43
- // Short-circuit the case of no tasks.
44
- if (n == 0) return;
45
-
46
- // Short-circuit the case of a single task.
47
- if (n == 1) {
48
- f(0);
49
- return;
50
- }
51
-
52
- // Heap-allocated state that manages concurrent execution of `f`.
53
- struct State {
54
- State(const Eigen::ThreadPoolDevice* intra_op_threadpool, F&& f)
55
- : intra_op_threadpool(intra_op_threadpool), f(std::forward<F>(f)) {}
56
-
57
- void Execute(std::shared_ptr<State> self, int64_t start, int64_t end) {
58
- while (end - start > 1) {
59
- uint64_t mid = (start + end) / 2;
60
- intra_op_threadpool->getPool()->Schedule(
61
- std::bind(&State::Execute, this, self, mid, end));
62
- end = mid;
63
- }
64
- f(start);
65
- }
66
-
67
- const Eigen::ThreadPoolDevice* intra_op_threadpool;
68
- F f;
69
- };
70
-
71
- auto s = std::make_shared<State>(intra_op_threadpool, std::forward<F>(f));
72
- s->Execute(std::move(s), 0, n);
73
- }
74
-
75
- } // namespace xla::cpu
76
-
77
- #endif // XLA_BACKENDS_CPU_RUNTIME_CONCURRENCY_H_