| // RUN: tf-mlir-translate -hlo-text-to-mlir-hlo %s -o - | FileCheck %s |
| |
| HloModule main |
| |
| // CHECK-LABEL: func @dynamic.update.slice.1(%arg0: tensor<4x4xf32>, %arg1: tensor<1x4xf32>, %arg2: tensor<f32>, %arg3: tensor<f32>) -> tensor<4x4xf32> { |
| %dynamic.update.slice.1 (Arg_0.1: f32[4, 4], Arg_1.2: f32[1, 4], Arg_2.3: f32[], Arg_3.4: f32[]) -> f32[4, 4] { |
| %Arg_0.1 = f32[4, 4] parameter(0) |
| %Arg_1.2 = f32[1, 4] parameter(1) |
| %Arg_2.3 = f32[] parameter(2) |
| %Arg_3.4 = f32[] parameter(3) |
| |
| // CHECK-NEXT: %0 = "xla_hlo.dynamic-update-slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<4x4xf32>, tensor<1x4xf32>, tensor<f32>, tensor<f32>) -> tensor<4x4xf32> |
| // CHECK-NEXT: return %0 : tensor<4x4xf32> |
| ROOT %dynamic-update-slice.5 = f32[4, 4] dynamic-update-slice(%Arg_0.1, %Arg_1.2, %Arg_2.3, %Arg_3.4) |
| } |
| |
| // CHECK-LABEL: func @main(%arg0: tensor<4xf32>, %arg1: tensor<2xf32>, %arg2: tensor<f32>) -> tensor<4xf32> |
| %dynamic.update.slice.2 (Arg_0.1: f32[4], Arg_1.2: f32[2], Arg_2.3: f32[]) -> f32[4] { |
| %Arg_0.1 = f32[4] parameter(0) |
| %Arg_1.2 = f32[2] parameter(1) |
| %Arg_2.3 = f32[] parameter(2) |
| |
| // CHECK-NEXT: %0 = "xla_hlo.dynamic-update-slice"(%arg0, %arg1, %arg2) : (tensor<4xf32>, tensor<2xf32>, tensor<f32>) -> tensor<4xf32> |
| // CHECK-NEXT: return %0 : tensor<4xf32> |
| ROOT %dynamic-update-slice.5 = f32[4] dynamic-update-slice(%Arg_0.1, %Arg_1.2, %Arg_2.3) |
| } |