tensorflow
diff --git a/RELEASE.md b/RELEASE.md
@@ -29,6 +29,7 @@ In `tensorflow/c/experimental/filesystem/filesystem_interface.h`, removed `TF_Tr
 
 * `BatchFunction Operator`
     * Adds the `num_warmup_batch_threads` op attribute to support a separate thread pool for processing warmup requests.
+    * Adds the `per_criticality_batch_timeout_micros` op attribute to support different batch timeouts for different criticality levels.
 * `TensorFlow API`
     * Exports `__new__` in public API golden files for subclasses of `tuple` (like `tf.io.FixedLenFeature`) to fix false positives during static type checking.>
 * `tf.data`

diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/passes/convert_tpu_model_to_cpu.td b/tensorflow/compiler/mlir/quantization/tensorflow/passes/convert_tpu_model_to_cpu.td
@@ -28,7 +28,7 @@ def CreateEmptyDictAttr : NativeCodeCall<"$_builder.getArrayAttr({})">;
 // because `TF_BatchFunctionOp` doesn't have the `CallOpInterface` trait.
 def ReplaceBatchFunctionOpToPartitionedCallOp : Pat<
   (TF_BatchFunctionOp:$src_op_res
-      $_, $_, $f, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_),
+      $_, $_, $f, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_),
   (TF_PartitionedCallOp
       (GetBatchFunctionOpArgOperands $src_op_res),
       /*arg_attrs=*/(CreateEmptyDictAttr),

diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_device.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_device.cc
@@ -151,7 +151,7 @@ LogicalResult ClusterFuncOp::verifySymbolUses(
 }
 
 void ClusterFuncOp::setCalleeFromCallable(mlir::CallInterfaceCallable callee) {
-  SymbolRefAttr calleeAttr = callee.get<SymbolRefAttr>();
+  SymbolRefAttr calleeAttr = cast<SymbolRefAttr>(callee);
   return setFuncAttr(cast<FlatSymbolRefAttr>(calleeAttr));
 }
 

diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
@@ -1141,6 +1141,7 @@ to be batched.}]>:$captured_tensors,
     DefaultValuedOptionalAttr<BoolAttr, "false">:$enable_large_batch_splitting,
     DefaultValuedOptionalAttr<BoolAttr, "false">:$enable_priority_aware_batch_scheduler,
     DefaultValuedOptionalAttr<BoolAttr, "false">:$enable_priority_aware_batch_scheduler_resplit,
+    DefaultValuedOptionalAttr<I64ArrayAttr, "{}">:$per_criticality_batch_timeout_micros,
     DefaultValuedOptionalAttr<BoolAttr, "false">:$enable_batching_task_lazy_cancellation,
     DefaultValuedOptionalAttr<I64Attr, "0">:$num_warmup_batch_threads
   );
@@ -12983,7 +12984,7 @@ def TF_RemoteCallOp : TF_Op<"RemoteCall", [CallOpInterface]> {
 
     /// Set the callee for this operation.
     void setCalleeFromCallable(CallInterfaceCallable callee) {
-      (*this)->setAttr("f", callee.get<SymbolRefAttr>());
+      (*this)->setAttr("f", ::llvm::cast<SymbolRefAttr>(callee));
     }
     Attribute removeArgAttrsAttr() { return nullptr; }
     Attribute removeResAttrsAttr() { return nullptr; }

diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td
@@ -2269,6 +2269,70 @@ def TF_XlaSparseDenseMatmulCustomCombinerOnTcWithCsrInputOp : TF_Op<"XlaSparseDe
   );
 }
 
+def TF_XlaSparseDenseMatmulPositionalWeightedWithCsrInputOp : TF_Op<"XlaSparseDenseMatmulPositionalWeightedWithCsrInput", [Pure]> {
+  let summary = "This op looks up the embedding vectors on SparseCores and performs the positional weighted combiner computation on SparseCores.";
+
+  let arguments = (ins
+    TF_Int32Tensor:$row_pointers,
+    TF_Int32Tensor:$sorted_sample_ids,
+    TF_Int32Tensor:$sorted_token_ids,
+    TF_Int32Tensor:$sorted_pos_ids,
+    TF_Float32Tensor:$embedding_table,
+    TF_Float32Tensor:$weights,
+
+    ConfinedAttr<I64Attr, [IntMinValue<0>]>:$input_size,
+    ConfinedAttr<I64Attr, [IntMinValue<0>]>:$num_weights,
+    ConfinedAttr<I64Attr, [IntMinValue<0>]>:$num_logical_devices,
+    OptionalAttr<F32Attr>:$quantization_config_low,
+    OptionalAttr<F32Attr>:$quantization_config_high,
+    OptionalAttr<I64Attr>:$quantization_config_num_buckets,
+
+    StrAttr:$table_name
+  );
+
+  let results = (outs
+    TF_Float32Tensor:$activations,
+    TF_Float32Tensor:$preserved_received_unique_vectors
+  );
+}
+
+def TF_XlaSparseDenseMatmulPositionalWeightedGradWithAdagradAndCsrInputOp : TF_Op<"XlaSparseDenseMatmulPositionalWeightedGradWithAdagradAndCsrInput", [AttrSizedOperandSegments, Pure]> {
+  let summary = "A XLA op which calculates the gradients and performs Adagrad update on the embedding table and custom optimizer update on positional weights and associated slot variables.";
+
+  let arguments = (ins
+    TF_Int32Tensor:$row_pointers,
+    TF_Int32Tensor:$sorted_sample_ids,
+    TF_Int32Tensor:$sorted_token_ids,
+    TF_Int32Tensor:$sorted_pos_ids,
+    TF_Float32Tensor:$activation_gradients,
+    TF_Float32Tensor:$preserved_received_unique_vectors,
+    TF_Float32Tensor:$preserved_weights,
+    TF_Float32Tensor:$embedding_table,
+    TF_Float32Tensor:$embedding_table_accumulator,
+    TF_Float32Tensor:$embedding_table_learning_rate,
+    Variadic<TF_Float32Tensor>:$weights,
+    Variadic<TF_Float32Tensor>:$hyperparameters,
+
+    F32Attr:$clip_weight_min,
+    F32Attr:$clip_weight_max,
+    ConfinedAttr<I64Attr, [IntMinValue<0>]>:$num_weights,
+    SymbolRefAttr:$custom_computation,
+    OptionalAttr<F32Attr>:$quantization_config_low,
+    OptionalAttr<F32Attr>:$quantization_config_high,
+    OptionalAttr<I64Attr>:$quantization_config_num_buckets,
+    StrAttr:$table_name
+  );
+
+  let results = (outs
+    TF_Float32Tensor:$updated_embedding_table,
+    TF_Float32Tensor:$updated_embedding_table_accumulator,
+    Variadic<TF_Float32Tensor>:$updated_weights
+  );
+
+  TF_DerivedOperandSizeAttr N = TF_DerivedOperandSizeAttr<10>;
+  TF_DerivedOperandSizeAttr M = TF_DerivedOperandSizeAttr<11>;
+}
+
 def TF_XlaSparseDenseMatmulGradWithSgdAndCsrInputOp : TF_Op<"XlaSparseDenseMatmulGradWithSgdAndCsrInput", [Pure]> {
   let summary = "";
 

diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.cc
@@ -3392,11 +3392,11 @@ LogicalResult LegacyCallOp::verifySymbolUses(
 void LegacyCallOp::setCalleeFromCallable(mlir::CallInterfaceCallable callee) {
   // Direct call.
   if (SymbolRefAttr fAttr = getFAttr()) {
-    SymbolRefAttr calleeAttr = callee.get<SymbolRefAttr>();
+    SymbolRefAttr calleeAttr = cast<SymbolRefAttr>(callee);
     return setFAttr(cast<FlatSymbolRefAttr>(calleeAttr));
   }
   // Indirect call, callee Value is the first operand.
-  return setOperand(0, callee.get<Value>());
+  return setOperand(0, cast<Value>(callee));
 }
 
 //===----------------------------------------------------------------------===//

diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.cc
@@ -601,11 +601,11 @@ static void SetPartitionCalleeFromCallable(CallOpClass op,
                                            mlir::CallInterfaceCallable callee) {
   // Direct call.
   if (SymbolRefAttr fAttr = op.getFAttr()) {
-    SymbolRefAttr calleeAttr = callee.get<SymbolRefAttr>();
+    SymbolRefAttr calleeAttr = cast<SymbolRefAttr>(callee);
     return op.setFAttr(cast<FlatSymbolRefAttr>(calleeAttr));
   }
   // Indirect call, callee Value is the first operand.
-  return op.setOperand(0, callee.get<Value>());
+  return op.setOperand(0, cast<Value>(callee));
 }
 
 void PartitionedCallOp::setCalleeFromCallable(

diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/cluster_ops_by_policy.cc b/tensorflow/compiler/mlir/tensorflow/transforms/cluster_ops_by_policy.cc
@@ -263,7 +263,7 @@ LogicalResult ClusteringState::VerifyDominanceProperty(
 
     // Block arguments do not really participate in clustering, they are only
     // used to connect independent operation using the same argument.
-    if (member.source.is<BlockArgument *>()) continue;
+    if (isa<BlockArgument*>(member.source)) continue;
 
     Operation *op = member.source.dyn_cast<Operation *>();
     assert(op && "member operation must be not null");

diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/region_control_flow_to_functional.cc b/tensorflow/compiler/mlir/tensorflow/transforms/region_control_flow_to_functional.cc
@@ -159,7 +159,7 @@ std::optional<StringRef> UnwrapSingleFunctionCall(Region& region) {
        llvm::zip(call.getResults(), yield.getOperands())) {
     if (ret != operand) return std::nullopt;
   }
-  SymbolRefAttr symbol = call.getCallableForCallee().get<SymbolRefAttr>();
+  SymbolRefAttr symbol = dyn_cast<SymbolRefAttr>(call.getCallableForCallee());
   if (!symbol) return std::nullopt;
   return symbol.getLeafReference();
 }

diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/shape_inference.cc b/tensorflow/compiler/mlir/tensorflow/transforms/shape_inference.cc
@@ -3130,7 +3130,7 @@ LogicalResult ShapeInference::TryToFold(Operation* op) {
     if ((attr = fold_result.dyn_cast<Attribute>())) {
       RecordValue(ValuePort(std::get<0>(result)), attr);
     } else {
-      auto value = fold_result.get<Value>();
+      auto value = cast<Value>(fold_result);
       assert(value.getType() == std::get<0>(result).getType() &&
              "folder produced value of incorrect type");
       if ((attr = ComputeOutputComponent(ValuePort(value)))) {

diff --git a/tensorflow/compiler/mlir/tf2xla/internal/passes/xla_cluster_formation.cc b/tensorflow/compiler/mlir/tf2xla/internal/passes/xla_cluster_formation.cc
@@ -107,7 +107,7 @@ mlir::LogicalResult EncapsulateFirstXlaCompilablePartitionedCalls(
   for (auto &pcall_op : outermost_pcall_ops) {
     auto call = llvm::cast<CallOpInterface>(pcall_op.getOperation());
     CallInterfaceCallable callable = call.getCallableForCallee();
-    auto sym = callable.get<mlir::SymbolRefAttr>();
+    auto sym = llvm::cast<mlir::SymbolRefAttr>(callable);
     EncapsulatePartitionedCall(pcall_op, sym.getRootReference());
   }
   // Partitioned calls are executed asynchronous. The calls outside of

diff --git a/tensorflow/compiler/mlir/tf2xla/transforms/legalization_op_config.cc b/tensorflow/compiler/mlir/tf2xla/transforms/legalization_op_config.cc
@@ -360,6 +360,9 @@ bool IsOpTypeAllowedTf2XlaFallback(const TypeID& type_id) {
         TypeID::get<TF::XlaSparseDenseMatmulGradWithSgdAndCsrInputOp>(),
         TypeID::get<TF::XlaSparseDenseMatmulWithCsrInputOp>(),
         TypeID::get<TF::XlaSparseDenseMatmulCustomCombinerOnTcWithCsrInputOp>(),
+        TypeID::get<TF::XlaSparseDenseMatmulPositionalWeightedWithCsrInputOp>(),
+        TypeID::get<
+            TF::XlaSparseDenseMatmulPositionalWeightedGradWithAdagradAndCsrInputOp>(),  // NOLINT
         TypeID::get<TF::XlaSparseDenseMatmulWithStaticBufferSizeOp>(),
         TypeID::get<
             TF::XlaSparseDenseMatmulGradWithAdagradAndStaticBufferSizeOp>(),

diff --git a/tensorflow/compiler/mlir/tf2xla/transforms/legalization_op_config_test.cc b/tensorflow/compiler/mlir/tf2xla/transforms/legalization_op_config_test.cc
@@ -83,7 +83,7 @@ TEST(LegalizationOpConfigTest, CountLoweringsSet) {
   // from MLIR to TF2XLA), these numbers should change. Or if TF Dialect adds
   // a new op, we should expect these to change too.
   EXPECT_EQ(mlir_lowering_count, 67);
-  EXPECT_EQ(tf2xla_fallback_count, 333);
+  EXPECT_EQ(tf2xla_fallback_count, 335);
   EXPECT_EQ(non_categorized_count, 437);
 }
 

diff --git a/tensorflow/compiler/mlir/tfr/ir/tfr_ops.cc b/tensorflow/compiler/mlir/tfr/ir/tfr_ops.cc
@@ -168,11 +168,11 @@ bool TFRType::classof(Type type) {
 void CallOp::setCalleeFromCallable(CallInterfaceCallable callee) {
   // Direct call.
   if (FlatSymbolRefAttr calleeAttr = getCalleeAttr()) {
-    auto symRef = callee.get<SymbolRefAttr>();
+    auto symRef = cast<SymbolRefAttr>(callee);
     return setCalleeAttr(cast<FlatSymbolRefAttr>(symRef));
   }
   // Indirect call, callee Value is the first operand.
-  return setOperand(0, callee.get<Value>());
+  return setOperand(0, cast<Value>(callee));
 }
 
 LogicalResult ConstantTensorOp::verify() {

diff --git a/tensorflow/compiler/mlir/tfrt/transforms/ifrt/BUILD b/tensorflow/compiler/mlir/tfrt/transforms/ifrt/BUILD
@@ -134,7 +134,6 @@ cc_library(
         "@llvm-project//mlir:Transforms",
         "@tsl//tsl/platform:protobuf",
         "@xla//xla:xla_data_proto_cc",
-        "@xla//xla/service:computation_placer_hdr",
     ],
 )
 
@@ -179,7 +178,7 @@ cc_library(
         "@xla//xla/hlo/translate/hlo_to_mhlo:hlo_to_mlir_hlo",
         "@xla//xla/pjrt:pjrt_compiler",
         "@xla//xla/python/ifrt",
-        "@xla//xla/service:computation_placer_hdr",
+        "@xla//xla/service:device_assignment",
         "@xla//xla/service:hlo_proto_cc",
         "@xla//xla/stream_executor:platform_manager",
         "@xla//xla/tsl/platform:errors",
@@ -245,7 +244,7 @@ tf_cc_test(
         "@xla//xla/python/ifrt:test_util",
         "@xla//xla/python/pjrt_ifrt",
         "@xla//xla/python/pjrt_ifrt:tfrt_cpu_client_test_lib",
-        "@xla//xla/service:computation_placer_hdr",
+        "@xla//xla/service:device_assignment",
         "@xla//xla/tsl/platform:statusor",
     ],
 )

diff --git a/tensorflow/compiler/mlir/tfrt/transforms/ifrt/sink_variable_as_named_array.cc b/tensorflow/compiler/mlir/tfrt/transforms/ifrt/sink_variable_as_named_array.cc
@@ -45,7 +45,6 @@ limitations under the License.
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
 #include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_constants.h"
-#include "xla/service/computation_placer.h"
 #include "xla/xla_data.pb.h"
 #include "tensorflow/core/platform/protobuf.h"  // IWYU pragma: keep
 #include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"

diff --git a/tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo.cc b/tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo.cc
@@ -56,7 +56,7 @@ limitations under the License.
 #include "xla/pjrt/pjrt_compiler.h"
 #include "xla/python/ifrt/client.h"
 #include "xla/python/ifrt/layout.h"
-#include "xla/service/computation_placer.h"
+#include "xla/service/device_assignment.h"
 #include "xla/shape.h"
 #include "xla/stream_executor/platform_manager.h"
 #include "xla/tsl/lib/strings/proto_serialization.h"

diff --git a/tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo_test.cc b/tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo_test.cc
@@ -46,7 +46,7 @@ limitations under the License.
 #include "xla/python/ifrt/mock.h"
 #include "xla/python/ifrt/test_util.h"
 #include "xla/python/pjrt_ifrt/pjrt_topology.h"
-#include "xla/service/computation_placer.h"
+#include "xla/service/device_assignment.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/tsl/lib/core/status_test_util.h"

diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/BUILD b/tensorflow/compiler/mlir/tools/kernel_gen/BUILD
@@ -209,7 +209,7 @@ cc_library(
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:string_view",
-        "@llvm-project//llvm:JITLink",
+        "@llvm-project//llvm:OrcJIT",
         "@llvm-project//llvm:OrcShared",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ExecutionEngine",

diff --git a/tensorflow/compiler/tests/image_ops_test.py b/tensorflow/compiler/tests/image_ops_test.py
@@ -877,6 +877,32 @@ def testNonAlignCorners3x2To6x4Batch2(self):
         out = sess.run(resized, {image: input_image[:, :, :, np.newaxis]})
         self.assertAllClose(expected[:, :, :, np.newaxis], out)
 
+  def testResizeBilinearNanPropagation(self):
+    # 2x2 image, 3 channels. One NaN at [0, 0, 1, 0] (row=0, col=1, channel=0)
+    input_data = [[
+        [[1.0, 2.0, 3.0], [float("nan"), 5.0, 6.0]],
+        [[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]],
+    ]]
+    cases = [
+        # (half_pixel_centers, align_corners, expected_nan_count)
+        (True, False, 9),
+        (False, False, 6),
+        (False, True, 9),
+    ]
+    for half_pixel_centers, align_corners, expected_nan_count in cases:
+      with self.session() as sess, self.test_scope():
+        image = array_ops.placeholder(np.float32)
+        resized = gen_image_ops.resize_bilinear(
+            image,
+            [4, 4],
+            half_pixel_centers=half_pixel_centers,
+            align_corners=align_corners,
+        )
+        out = sess.run(resized, {image: np.array(input_data, dtype=np.float32)})
+        self.assertEqual(np.isnan(out[0, :, :, 0]).sum(), expected_nan_count)
+        self.assertEqual(np.isnan(out[0, :, :, 1]).sum(), 0)
+        self.assertEqual(np.isnan(out[0, :, :, 2]).sum(), 0)
+
 
 class ResizeBilinearGradHalfPixelCentersTest(
     parameterized.TestCase, xla_test.XLATestCase

diff --git a/tensorflow/compiler/tests/scan_ops_test.py b/tensorflow/compiler/tests/scan_ops_test.py
@@ -124,6 +124,12 @@ def test6D(self):
       for axis in range(-6, 6, 3):
         self._compareAll(x, axis)
 
+  def testComplex(self):
+    x = np.array([[1 + 2j, 3 + 4j], [5 + 6j, 7 + 8j]])
+    for dtype in self.complex_types:
+      for axis in (-2, -1, 0, 1):
+        self._compareAll(x.astype(dtype), axis)
+
   def testMixedPrecision(self):
     with self.session(), self.test_scope():
       y = math_ops.cumsum(
@@ -296,6 +302,12 @@ def test6D(self):
       for axis in range(-6, 6, 3):
         self._compareAll(x, axis)
 
+  def testComplex(self):
+    x = np.array([[1 + 2j, 3 + 4j], [5 + 6j, 7 + 8j]])
+    for dtype in self.complex_types:
+      for axis in (-2, -1, 0, 1):
+        self._compareAll(x.astype(dtype), axis)
+
   @test_util.disable_mlir_bridge("Error handling")
   def testInvalidAxis(self):
     x = np.arange(0, 10).reshape([2, 5]).astype(np.float32)

diff --git a/tensorflow/compiler/tests/tensor_array_ops_test.py b/tensorflow/compiler/tests/tensor_array_ops_test.py
@@ -454,8 +454,11 @@ def fn():
             infer_shape=False)
         return ta.split([1.0, 2.0, 3.0], [1, 2, 3]).flow
 
-      with self.assertRaisesOpError(
-          r"lengths must be equal: 1 vs. 2"):
+      with self.assertRaisesRegex(
+          ValueError,
+          r"Expected sum of lengths to be equal to values.shape\[0\], "
+          r"but sum of lengths is 6 and value's shape is: \[3\]",
+      ):
         xla.compile(fn)[0].eval()
 
       def fn():
@@ -466,8 +469,10 @@ def fn():
             infer_shape=False)
         return ta.split(1.0, [1]).flow
 
-      with self.assertRaisesOpError(
-          r"value must have rank >= 1"):
+      with self.assertRaisesRegex(
+          ValueError,
+          r"Expected value to be at least a vector, but received shape: \[\]",
+      ):
         xla.compile(fn)[0].eval()
 
       def fn():

diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD
@@ -688,7 +688,7 @@ cc_library(
         "@xla//xla/hlo/builder:xla_computation",
         "@xla//xla/hlo/ir:hlo",
         "@xla//xla/hlo/translate/mhlo_to_hlo:layout_util",
-        "@xla//xla/service:computation_placer_hdr",
+        "@xla//xla/service:device_assignment",
         "@xla//xla/service:hlo_proto_cc",
         "@xla//xla/service/spmd/shardy:constants",
         "@xla//xla/service/spmd/shardy:utils",
@@ -895,7 +895,7 @@ cc_library(
         "@xla//xla/hlo/builder:xla_computation",
         "@xla//xla/hlo/ir:hlo",
         "@xla//xla/hlo/translate/mhlo_to_hlo:layout_util",
-        "@xla//xla/service:computation_placer_hdr",
+        "@xla//xla/service:device_assignment",
         "@xla//xla/service/gpu:gpu_executable_run_options",
         "@xla//xla/stream_executor:stream",
         "@xla//xla/stream_executor:stream_executor_h",