[pull] master from tensorflow:master

RELEASE.md

@@ -46,6 +46,14 @@ In `tensorflow/c/experimental/filesystem/filesystem_interface.h`, removed `TF_Tr
         previously failed with a lookup error because the `SoftsignGrad`
         backward op had no registered Python gradient.
 
+*   `tf.math.reciprocal`
+
+    *   Constrains the XLA registration of `Reciprocal` and `Inv` to the types
+        that have a device kernel, so `jit_compile=True` no longer silently
+        accepts the integer inputs that eager execution and autoclustering
+        reject. Fixes
+        [#126414](https://github.com/tensorflow/tensorflow/issues/126414).
+
 
 *   `tf.experimental.numpy`
 

ci/official/envs/linux_x86_cuda

@@ -14,7 +14,7 @@
 # ==============================================================================
 source ci/official/envs/linux_x86
 export TF_FORCE_GPU_ALLOW_GROWTH=true
-TFCI_BAZEL_COMMON_ARGS="--repo_env=HERMETIC_PYTHON_VERSION=$TFCI_PYTHON_VERSION --repo_env=USE_PYWRAP_RULES=True --config release_gpu_linux --test_env=TF_FORCE_GPU_ALLOW_GROWTH=true"
+TFCI_BAZEL_COMMON_ARGS="--repo_env=HERMETIC_PYTHON_VERSION=$TFCI_PYTHON_VERSION --repo_env=USE_PYWRAP_RULES=True --config release_gpu_linux --test_env=TF_FORCE_GPU_ALLOW_GROWTH=true --local_test_jobs=16"
 TFCI_BAZEL_HERMETIC_CUDA_UMD_ENABLE=1
 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_cuda
 TFCI_BUILD_PIP_PACKAGE_WHEEL_NAME_ARG="--repo_env=WHEEL_NAME=tensorflow"

ci/official/envs/linux_x86_cuda13_nvcc

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 source ci/official/envs/linux_x86
-TFCI_BAZEL_COMMON_ARGS="--repo_env=HERMETIC_PYTHON_VERSION=$TFCI_PYTHON_VERSION --repo_env=USE_PYWRAP_RULES=True --config release_gpu_linux --config=cuda_nvcc --config=cuda13_version"
+TFCI_BAZEL_COMMON_ARGS="--repo_env=HERMETIC_PYTHON_VERSION=$TFCI_PYTHON_VERSION --repo_env=USE_PYWRAP_RULES=True --config release_gpu_linux --config=cuda_nvcc --config=cuda13_version --local_test_jobs=16"
 TFCI_BAZEL_HERMETIC_CUDA_UMD_ENABLE=1
 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_cuda_13_nvcc
 TFCI_BUILD_PIP_PACKAGE_WHEEL_NAME_ARG="--repo_env=WHEEL_NAME=tensorflow_cuda13"

requirements_lock_3_10.txt

@@ -1077,9 +1077,9 @@ zstandard==0.25.0 \
     # via -r ci/official/requirements_updater/requirements.in
 
 # The following packages are considered to be unsafe in a requirements file:
-setuptools==78.1.1 \
-    --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \
-    --hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d
+setuptools==83.0.0 \
+    --hash=sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef \
+    --hash=sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3
     # via
     #   -r ci/official/requirements_updater/requirements.in
     #   tb-nightly

requirements_lock_3_11.txt

@@ -1076,9 +1076,9 @@ zstandard==0.25.0 \
     # via -r ci/official/requirements_updater/requirements.in
 
 # The following packages are considered to be unsafe in a requirements file:
-setuptools==78.1.1 \
-    --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \
-    --hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d
+setuptools==83.0.0 \
+    --hash=sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef \
+    --hash=sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3
     # via
     #   -r ci/official/requirements_updater/requirements.in
     #   tb-nightly

requirements_lock_3_12.txt

@@ -1076,9 +1076,9 @@ zstandard==0.25.0 \
     # via -r ci/official/requirements_updater/requirements.in
 
 # The following packages are considered to be unsafe in a requirements file:
-setuptools==78.1.1 \
-    --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \
-    --hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d
+setuptools==83.0.0 \
+    --hash=sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef \
+    --hash=sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3
     # via
     #   -r ci/official/requirements_updater/requirements.in
     #   tb-nightly

requirements_lock_3_13.txt

@@ -1083,9 +1083,9 @@ zstandard==0.25.0 \
     # via -r ci/official/requirements_updater/requirements.in
 
 # The following packages are considered to be unsafe in a requirements file:
-setuptools==78.1.1 \
-    --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \
-    --hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d
+setuptools==83.0.0 \
+    --hash=sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef \
+    --hash=sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3
     # via
     #   -r ci/official/requirements_updater/requirements.in
     #   tb-nightly

requirements_lock_3_14.txt

@@ -1317,9 +1317,9 @@ zstandard==0.25.0 \
     # via -r ci/official/requirements_updater/requirements.in
 
 # The following packages are considered to be unsafe in a requirements file:
-setuptools==78.1.1 \
-    --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \
-    --hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d
+setuptools==83.0.0 \
+    --hash=sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef \
+    --hash=sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3
     # via
     #   -r ci/official/requirements_updater/requirements.in
     #   tb-nightly

requirements_lock_3_14_freethreaded.txt

@@ -1318,9 +1318,9 @@ zstandard==0.25.0 \
     # via -r ci/official/requirements_updater/requirements.in
 
 # The following packages are considered to be unsafe in a requirements file:
-setuptools==78.1.1 \
-    --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \
-    --hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d
+setuptools==83.0.0 \
+    --hash=sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef \
+    --hash=sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3
     # via
     #   -r ci/official/requirements_updater/requirements.in
     #   tb-nightly

tensorflow/compiler/mlir/lite/tests/optimize.mlir

@@ -4067,6 +4067,46 @@ func.func @gelu_approximate1_with_mul1(%arg0: tensor<3xf32>) -> tensor<3xf32> {
 // CHECK: "tfl.gelu"(%arg0) <{approximate = true}> : (tensor<3xf32>) -> tensor<3xf32>
 }
 
+func.func @gelu_approximate_with_mul3(%arg0: tensor<3xf32>) -> tensor<3xf32> {
+  %cst = arith.constant dense<0.797884583> : tensor<f32>
+  %cst_0 = arith.constant dense<5.000000e-01> : tensor<f32>
+  %cst_1 = arith.constant dense<1.000000e+00> : tensor<f32>
+  %cst_3 = arith.constant dense<4.471500e-02> : tensor<f32>
+  %0 = "tfl.mul"(%arg0, %cst_3) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<f32>) -> tensor<3xf32>
+  %1 = "tfl.mul"(%0, %arg0) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32>
+  %2 = "tfl.mul"(%1, %arg0) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32>
+  %3 = "tfl.add"(%arg0, %2) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32>
+  %4 = "tfl.mul"(%3, %cst) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<f32>) -> tensor<3xf32>
+  %5 = "tfl.tanh"(%4) : (tensor<3xf32>) -> tensor<3xf32>
+  %6 = "tfl.add"(%5, %cst_1) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<f32>) -> tensor<3xf32>
+  %7 = "tfl.mul"(%arg0, %cst_0) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<f32>) -> tensor<3xf32>
+  %8 = "tfl.mul"(%7, %6) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32>
+  func.return %8 : tensor<3xf32>
+
+// CHECK-LABEL:gelu_approximate
+// CHECK: "tfl.gelu"(%arg0) <{approximate = true}> : (tensor<3xf32>) -> tensor<3xf32>
+}
+
+func.func @gelu_approximate1_with_mul3(%arg0: tensor<3xf32>) -> tensor<3xf32> {
+  %cst = arith.constant dense<0.797884583> : tensor<f32>
+  %cst_0 = arith.constant dense<5.000000e-01> : tensor<f32>
+  %cst_1 = arith.constant dense<1.000000e+00> : tensor<f32>
+  %cst_3 = arith.constant dense<4.471500e-02> : tensor<f32>
+  %0 = "tfl.mul"(%arg0, %cst_3) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<f32>) -> tensor<3xf32>
+  %1 = "tfl.mul"(%0, %arg0) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32>
+  %2 = "tfl.mul"(%1, %arg0) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32>
+  %3 = "tfl.add"(%arg0, %2) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32>
+  %4 = "tfl.mul"(%3, %cst) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<f32>) -> tensor<3xf32>
+  %5 = "tfl.tanh"(%4) : (tensor<3xf32>) -> tensor<3xf32>
+  %6 = "tfl.add"(%5, %cst_1) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<f32>) -> tensor<3xf32>
+  %7 = "tfl.mul"(%6, %cst_0) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<f32>) -> tensor<3xf32>
+  %8 = "tfl.mul"(%arg0, %7) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32>
+  func.return %8 : tensor<3xf32>
+
+// CHECK-LABEL:gelu_approximate
+// CHECK: "tfl.gelu"(%arg0) <{approximate = true}> : (tensor<3xf32>) -> tensor<3xf32>
+}
+
 func.func @gelu_approximate_no_match(%arg0: tensor<3xf32>) -> tensor<3xf32> {
   %cst = arith.constant dense<0.797884583> : tensor<f32>
   %cst_0 = arith.constant dense<5.000000e-01> : tensor<f32>

tensorflow/compiler/mlir/lite/transforms/optimize_patterns.td

@@ -1566,6 +1566,41 @@ def MatchGeluApproximate_Mul2 : Pat<
    (HasOneUse $sqr_out),
   ]>;
 
+// Alternate pattern for GeluApproximate to match mul(mul(mul(x, coeff), x), x),
+// the cube shape produced by the left-associative spelling
+// `0.044715 * x * x * x` (common in Python model code), where the coefficient
+// is folded into the innermost mul so no bare x^3 subterm exists.
+//   0.5 * x * ( 1 + tanh( sqrt_2dPi  * ( x + mul(mul(mul(x, 0.044715), x), x) ) ) )
+def MatchGeluApproximate_Mul3 : Pat<
+  (TFL_MulOp
+   (TFL_MulOp:$mul_out $arg0, (Arith_ConstantOp F32ElementsAttr:$Cst_1_2), TFL_AF_None),
+   (TFL_AddOp:$add_out
+    (TFL_TanhOp:$tanh_out
+     (TFL_MulOp:$mul_out1
+      (TFL_AddOp:$add_out1 $arg0,
+       (TFL_MulOp:$mul_out2
+        (TFL_MulOp:$mul_out3
+         (TFL_MulOp:$mul_out4 $arg0,
+          (Arith_ConstantOp F32ElementsAttr:$Coeff), TFL_AF_None),
+         $arg0, TFL_AF_None),
+        $arg0, TFL_AF_None), TFL_AF_None),
+      (Arith_ConstantOp F32ElementsAttr:$Cst_sqrt_2dPi), TFL_AF_None)),
+    (Arith_ConstantOp F32ElementsAttr:$Cst_1), TFL_AF_None), TFL_AF_None),
+  (TFL_GeluOp $arg0, ConstBoolAttrTrue),
+  [(FloatValueEquals<"0.5"> $Cst_1_2),
+   (FloatValueEquals<"1"> $Cst_1),
+   (FloatValueEquals<"0.797884583"> $Cst_sqrt_2dPi),
+   (FloatValueEquals<"0.044715"> $Coeff),
+   (HasOneUse $mul_out),
+   (HasOneUse $add_out),
+   (HasOneUse $tanh_out),
+   (HasOneUse $mul_out1),
+   (HasOneUse $add_out1),
+   (HasOneUse $mul_out2),
+   (HasOneUse $mul_out3),
+   (HasOneUse $mul_out4),
+  ]>;
+
 // Alternate pattern for GeluApproximate (see different order for mul), replaces
 //   x * ( 0.5 * ( 1 + tanh( sqrt_2dPi  * ( x + 0.044715 * pow( x, 3 ) ) ) ) )
 def MatchGeluApproximate1 : Pat<
@@ -1656,6 +1691,38 @@ def MatchGeluApproximate1_Mul2 : Pat<
    (HasOneUse $sqr_out),
   ]>;
 
+// Alternate pattern for GeluApproximate1 to match mul(mul(mul(x, coeff), x), x).
+//   x * ( 0.5 * ( 1 + tanh( sqrt_2dPi  * ( x + mul(mul(mul(x, 0.044715), x), x) ) ) ) )
+def MatchGeluApproximate1_Mul3 : Pat<
+  (TFL_MulOp $arg0,
+   (TFL_MulOp:$mul_out
+    (TFL_AddOp:$add_out
+     (TFL_TanhOp:$tanh_out
+      (TFL_MulOp:$mul_out1
+       (TFL_AddOp:$add_out1 $arg0,
+        (TFL_MulOp:$mul_out2
+         (TFL_MulOp:$mul_out3
+          (TFL_MulOp:$mul_out4 $arg0,
+           (Arith_ConstantOp F32ElementsAttr:$Coeff), TFL_AF_None),
+          $arg0, TFL_AF_None),
+         $arg0, TFL_AF_None), TFL_AF_None),
+       (Arith_ConstantOp F32ElementsAttr:$Cst_sqrt_2dPi), TFL_AF_None)),
+     (Arith_ConstantOp F32ElementsAttr:$Cst_1), TFL_AF_None), (Arith_ConstantOp F32ElementsAttr:$Cst_1_2), TFL_AF_None), TFL_AF_None),
+  (TFL_GeluOp $arg0, ConstBoolAttrTrue),
+  [(FloatValueEquals<"0.5"> $Cst_1_2),
+   (FloatValueEquals<"1"> $Cst_1),
+   (FloatValueEquals<"0.797884583"> $Cst_sqrt_2dPi),
+   (FloatValueEquals<"0.044715"> $Coeff),
+   (HasOneUse $mul_out),
+   (HasOneUse $add_out),
+   (HasOneUse $tanh_out),
+   (HasOneUse $mul_out1),
+   (HasOneUse $add_out1),
+   (HasOneUse $mul_out2),
+   (HasOneUse $mul_out3),
+   (HasOneUse $mul_out4),
+  ]>;
+
 // For Gelu, replaces
 //   0.5 * x * ( 1 + erf( x * sqrt_1_2 ) )
 def MatchGelu : Pat<

tensorflow/compiler/tests/BUILD

@@ -1811,11 +1811,13 @@ tf_xla_py_strict_test(
     ],
     deps = [
         ":xla_test",
+        "//tensorflow/python/eager:def_function",
         "//tensorflow/python/framework:constant_op",
         "//tensorflow/python/framework:dtypes",
         "//tensorflow/python/framework:errors",
         "//tensorflow/python/ops:array_ops",
         "//tensorflow/python/ops:list_ops",
+        "//tensorflow/python/ops:map_fn",
         "//tensorflow/python/platform:client_testlib",
         "//third_party/py/numpy",
         "@absl_py//absl/testing:parameterized",

tensorflow/compiler/tests/slice_ops_test.py

@@ -50,6 +50,28 @@ def testZeroSlice(self):
 
         self.assertAllEqual([], result)
 
+  def testSliceOfDynamicDimension(self):
+    # Regression test for GitHub issue 110789. The input's leading dimension
+    # is dynamic because it comes from where, and a partial slice of such a
+    # dimension returned wrong output sizes, and corrupted the heap on some
+    # platforms, instead of taking the requested slice.
+    with self.session():
+      i = array_ops.placeholder(dtypes.int64, shape=[2, 2, 5])
+      with self.test_scope():
+        indices = array_ops.where(math_ops.not_equal(i, 0))
+        sliced = array_ops.slice(indices, [0, 0], [2, 1])
+        empty = array_ops.slice(indices, [0, 0], [0, 1])
+      params = {
+          i: [
+              [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]],
+              [[10, 11, 12, 13, 14], [15, 16, 17, 18, 19]],
+          ],
+      }
+      # The first two nonzero elements sit at indices (0, 0, 1) and
+      # (0, 0, 2), so the slice holds the first coordinate of each.
+      self.assertAllEqual([[0], [0]], sliced.eval(feed_dict=params))
+      self.assertAllEqual((0, 1), empty.eval(feed_dict=params).shape)
+
   def test3D(self):
     for dtype in self.numeric_types:
       with self.session():

tensorflow/compiler/tests/tensor_list_ops_test.py

@@ -19,17 +19,47 @@
 
 from absl.testing import parameterized
 import numpy as np
+
 from tensorflow.compiler.tests import xla_test
+from tensorflow.python.eager import def_function
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import list_ops
+from tensorflow.python.ops import map_fn
 from tensorflow.python.platform import test
 
 
 class ListOpsTest(parameterized.TestCase, xla_test.XLATestCase):
 
+  def testGetItemFromEmptyList(self):
+    # Regression test for GitHub issue 109648. Reading from a statically
+    # empty list appears in code that never runs, such as the body of a
+    # while loop with a zero trip count, but XLA compiles that code anyway
+    # and used to reject the read at compile time. It now yields zeros of
+    # the element shape.
+    with self.session() as sess, self.test_scope():
+      l = list_ops.tensor_list_reserve(
+          element_shape=[2], element_dtype=dtypes.float32, num_elements=0
+      )
+      e = list_ops.tensor_list_get_item(l, 0, element_dtype=dtypes.float32)
+      self.assertAllEqual(sess.run(e), [0.0, 0.0])
+
+  def testMapFnOverEmptyTensor(self):
+    # End to end case for GitHub issue 109648: map_fn over a zero length
+    # tensor compiles its loop body even though it never runs, and the
+    # TensorListGetItem in that body used to fail compilation. The whole
+    # function is jit-compiled so the list stays inside XLA rather than
+    # crossing the XLA/TF boundary at the unstack and stack ops.
+    @def_function.function(jit_compile=True)
+    def f(x):
+      return map_fn.map_fn(lambda t: t + 1.0, x)
+
+    with self.session() as sess:
+      x = array_ops.zeros([0], dtype=dtypes.float32)
+      self.assertAllEqual(sess.run(f(x)).shape, (0,))
+
   def testElementShape(self):
     with self.session() as sess, self.test_scope():
       dim = array_ops.placeholder(dtypes.int32)

tensorflow/compiler/tf2xla/BUILD

@@ -1643,11 +1643,13 @@ tf_cc_test(
     srcs = ["xla_op_registry_test.cc"],
     deps = [
         ":xla_compiler",
+        "//tensorflow/compiler/tf2xla/kernels:xla_ops",
         "//tensorflow/core:framework",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core:test",
         "//tensorflow/core:test_main",
         "@com_google_absl//absl/log",
+        "@com_google_googletest//:gtest",
     ],
 )
 

tensorflow/compiler/tf2xla/kernels/BUILD

@@ -1587,6 +1587,7 @@ tf_kernel_library(
         "//tensorflow/core:lib",
         "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/types:span",
+        "@xla//xla:shape_util",
         "@xla//xla:xla_data_proto_cc",
         "@xla//xla/hlo/builder:value_inference",
         "@xla//xla/hlo/builder:xla_builder",
@@ -1873,6 +1874,7 @@ tf_kernel_library(
         "//tensorflow/compiler/tf2xla:xla_resource",
         "//tensorflow/compiler/tf2xla/ops:xla_ops",
         "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
         "@com_google_absl//absl/status:statusor",
         "@xla//xla:xla_data_proto_cc",
         "@xla//xla/hlo/builder:xla_builder",

tensorflow/compiler/tf2xla/kernels/slice_op.cc

@@ -26,6 +26,7 @@ limitations under the License.
 #include "xla/hlo/builder/lib/dynamic_shaped_ops.h"
 #include "xla/hlo/builder/value_inference.h"
 #include "xla/hlo/builder/xla_builder.h"
+#include "xla/shape.h"
 #include "xla/xla_data.pb.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/op_requires.h"
@@ -104,6 +105,30 @@ class SliceOp : public XlaOpKernel {
       }
       std::vector<int64_t> strides(begin.size(), 1);
       auto slice = xla::Slice(ctx->Input(0), begin, limits, strides);
+
+      // If the input has dynamic dimensions, the dynamic size that XLA
+      // infers for a partial slice of such a dimension does not implement
+      // the slice semantics, so set the output size explicitly. A dimension
+      // holding `dynamic_size` valid elements contributes
+      // clamp(dynamic_size - begin, 0, size) elements to the slice.
+      auto input_xla_shape = ctx->builder()->GetShape(ctx->Input(0));
+      OP_REQUIRES_OK(ctx, input_xla_shape.status());
+      for (int64_t i = 0; i < input_dims; ++i) {
+        if (input_xla_shape->is_dynamic_dimension(i)) {
+          xla::XlaOp input_dynamic_size =
+              xla::GetDimensionSize(ctx->Input(0), i);
+          xla::XlaOp begin_size = xla::ConstantR0<int32_t>(
+              ctx->builder(), static_cast<int32_t>(begin[i]));
+          xla::XlaOp requested_size = xla::ConstantR0<int32_t>(
+              ctx->builder(), static_cast<int32_t>(wrapped_size[i]));
+          xla::XlaOp output_size =
+              xla::Clamp(xla::ConstantR0<int32_t>(ctx->builder(), 0),
+                         input_dynamic_size - begin_size, requested_size);
+          slice = xla::RemoveDynamicDimension(slice, i);
+          slice = xla::SetDimensionSize(slice, output_size, i);
+        }
+      }
+
       // Check for slice on dynamic dimensions.
       std::vector<bool> size_is_dynamic;
       OP_REQUIRES_OK(