Files
tensorflow--tensorflow/tensorflow/compiler/mlir/lite/tests/const-fold.mlir
T
wehub-resource-sync 8a852e4b4e
cffconvert / validate (push) Has been skipped
License Check / license-check (push) Failing after 2s
chore: import upstream snapshot with attribution
2026-07-13 12:14:16 +08:00

1234 lines
52 KiB
MLIR

// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: litert-opt %s -canonicalize | env FILECHECK_OPTS="" FileCheck %s
// CHECK-LABEL: @elementwise_unary_ops
func.func @elementwise_unary_ops() -> (tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>) {
%0 = arith.constant dense<-1.0> : tensor<f32>
%1 = arith.constant dense<1.0> : tensor<f32>
%2 = arith.constant dense<1.0> : tensor<f32>
%3 = arith.constant dense<1.0> : tensor<f32>
%4 = arith.constant dense<4.0> : tensor<f32>
%5 = arith.constant dense<4.0> : tensor<f32>
%6 = arith.constant dense<2.0> : tensor<f32>
%one = arith.constant dense<1.0> : tensor<f32>
// CHECK-DAG: [[cst0:%.*]] = arith.constant dense<1.000000e+00> : tensor<f32>
// CHECK-DAG: [[cst1:%.*]] = arith.constant dense<0.841470957> : tensor<f32>
// CHECK-DAG: [[cst2:%.*]] = arith.constant dense<0.540302277> : tensor<f32>
// CHECK-DAG: [[cst3:%.*]] = arith.constant dense<0.000000e+00> : tensor<f32>
// CHECK-DAG: [[cst4:%.*]] = arith.constant dense<2.000000e+00> : tensor<f32>
// CHECK-DAG: [[cst5:%.*]] = arith.constant dense<5.000000e-01> : tensor<f32>
// CHECK-DAG: [[cst6:%.*]] = arith.constant dense<4.000000e+00> : tensor<f32>
// CHECK-DAG: [[cst7:%.*]] = arith.constant dense<0.761594176> : tensor<f32>
// CHECK-DAG: [[cst8:%.*]] = arith.constant dense<0.73105859{{[78]}}> : tensor<f32>
// CHECK: return [[cst0]], [[cst1]], [[cst2]], [[cst3]], [[cst4]], [[cst5]], [[cst6]], [[cst7]], [[cst8]]
%7 = "tfl.abs"(%0) : (tensor<f32>) -> tensor<f32>
%8 = "tfl.sin"(%1) : (tensor<f32>) -> tensor<f32>
%9 = "tfl.cos"(%2) : (tensor<f32>) -> tensor<f32>
%10 = "tfl.log"(%3) : (tensor<f32>) -> tensor<f32>
%11 = "tfl.sqrt"(%4) : (tensor<f32>) -> tensor<f32>
%12 = "tfl.rsqrt"(%5) : (tensor<f32>) -> tensor<f32>
%13 = "tfl.square"(%6) : (tensor<f32>) -> tensor<f32>
%14 = "tfl.tanh"(%one) : (tensor<f32>) -> tensor<f32>
%15 = "tfl.logistic"(%one) : (tensor<f32>) -> tensor<f32>
func.return %7, %8, %9, %10, %11, %12, %13, %14, %15 : tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>
}
// CHECK-LABEL: @cumsum_chained_pipeline
func.func @cumsum_chained_pipeline() -> tensor<3xf32> {
%input = arith.constant dense<[1.0, 2.0, 3.0]> : tensor<3xf32>
%axis = arith.constant dense<0> : tensor<i32>
%cst_two = arith.constant dense<2.0> : tensor<3xf32>
%one_tensor = arith.constant dense<1.0> : tensor<3xf32>
// CHECK: %[[PIPELINE_RES:.*]] = arith.constant dense<[0.982013761, 0.999664664, 0.999999165]> : tensor<3xf32>
// CHECK: return %[[PIPELINE_RES]]
%0 = "tfl.cumsum"(%input, %axis) {exclusive = false, reverse = false} : (tensor<3xf32>, tensor<i32>) -> tensor<3xf32>
%add = "tfl.add"(%0, %cst_two) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32>
%sub = "tfl.sub"(%add, %one_tensor) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32>
%mul = "tfl.mul"(%sub, %cst_two) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32>
%res_pipeline = "tfl.logistic"(%mul) : (tensor<3xf32>) -> tensor<3xf32>
func.return %res_pipeline : tensor<3xf32>
}
// CHECK-LABEL: @rank
func.func @rank() -> tensor<1xi32> {
%cst = arith.constant dense<[[1], [2]]> : tensor<2x1xi32>
// CHECK: %[[CST:.*]] = arith.constant dense<2> : tensor<1xi32>
// CHECK: return %[[CST]]
%0 = "tfl.rank"(%cst) : (tensor<2x1xi32>) -> tensor<1xi32>
func.return %0 : tensor<1xi32>
}
// CHECK-LABEL: @rank_input_known_rank
func.func @rank_input_known_rank(%arg0 : tensor<2x1xi32>) -> tensor<1xi32> {
// CHECK: %[[CST:.*]] = arith.constant dense<2> : tensor<1xi32>
// CHECK: return %[[CST]]
%0 = "tfl.rank"(%arg0) : (tensor<2x1xi32>) -> tensor<1xi32>
func.return %0 : tensor<1xi32>
}
// CHECK-LABEL: @reshape_dynamic_output
func.func @reshape_dynamic_output() -> tensor<?xi32> {
%input = arith.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>
%shape = arith.constant dense<[4]> : tensor<1xi32>
// CHECK: %[[CST:.*]] = "tfl.pseudo_const"() <{value = dense<[1, 2, 3, 4]> : tensor<4xi32>}> : () -> tensor<?xi32>
// CHECK: return %[[CST]]
%0 = "tfl.reshape"(%input, %shape) : (tensor<2x2xi32>, tensor<1xi32>) -> tensor<?xi32>
func.return %0 : tensor<?xi32>
}
// CHECK-LABEL: @pseudo_const
func.func @pseudo_const() -> tensor<i32> {
// CHECK: %[[CST:.*]] = arith.constant dense<1> : tensor<i32>
// CHECK: return %[[CST]]
%0 = "tfl.pseudo_const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
func.return %0 : tensor<i32>
}
// CHECK-LABEL: @range_int
func.func @range_int() -> tensor<?xi32> {
%cst = arith.constant dense<0> : tensor<i32>
%cst_1 = arith.constant dense<4> : tensor<i32>
%cst_2 = arith.constant dense<1> : tensor<i32>
// CHECK: %[[CST:.*]] = "tfl.pseudo_const"() <{value = dense<[0, 1, 2, 3]> : tensor<4xi32>}> : () -> tensor<?xi32>
// CHECK: return %[[CST]]
%0 = "tfl.range"(%cst, %cst_1, %cst_2) : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<?xi32>
func.return %0 : tensor<?xi32>
}
// CHECK-LABEL: @range_float
func.func @range_float() -> tensor<?xf32> {
%cst = arith.constant dense<0.0> : tensor<f32>
%cst_1 = arith.constant dense<4.0> : tensor<f32>
%cst_2 = arith.constant dense<1.0> : tensor<f32>
// CHECK: %[[CST:.*]] = "tfl.pseudo_const"() <{value = dense<[0.000000e+00, 1.000000e+00, 2.000000e+00, 3.000000e+00]> : tensor<4xf32>}> : () -> tensor<?xf32>
// CHECK: return %[[CST]]
%0 = "tfl.range"(%cst, %cst_1, %cst_2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<?xf32>
func.return %0 : tensor<?xf32>
}
// CHECK-LABEL: @range_float_neg_delta
func.func @range_float_neg_delta() -> tensor<?xf32> {
%cst = arith.constant dense<0.0> : tensor<f32>
%cst_1 = arith.constant dense<-4.0> : tensor<f32>
%cst_2 = arith.constant dense<-1.0> : tensor<f32>
// CHECK: %[[CST:.*]] = "tfl.pseudo_const"() <{value = dense<[0.000000e+00, -1.000000e+00, -2.000000e+00, -3.000000e+00]> : tensor<4xf32>}> : () -> tensor<?xf32>
// CHECK: return %[[CST]]
%0 = "tfl.range"(%cst, %cst_1, %cst_2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<?xf32>
func.return %0 : tensor<?xf32>
}
// CHECK-LABEL: @range_float_nonzero_base
func.func @range_float_nonzero_base() -> tensor<?xf32> {
%cst = arith.constant dense<2.0> : tensor<f32>
%cst_1 = arith.constant dense<7.0> : tensor<f32>
%cst_2 = arith.constant dense<1.5> : tensor<f32>
// CHECK: %[[CST:.*]] = "tfl.pseudo_const"() <{value = dense<[2.000000e+00, 3.500000e+00, 5.000000e+00, 6.500000e+00]> : tensor<4xf32>}> : () -> tensor<?xf32>
// CHECK: return %[[CST]]
%0 = "tfl.range"(%cst, %cst_1, %cst_2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<?xf32>
func.return %0 : tensor<?xf32>
}
// CHECK-LABEL: @transpose_dynamic
func.func @transpose_dynamic() -> tensor<?xi32> {
%cst = arith.constant dense<[1, 2, 3]> : tensor<3xi32>
%cst_perm = arith.constant dense<0> : tensor<1xi32>
// CHECK: %[[CST:.*]] = "tfl.pseudo_const"() <{value = dense<{{\[}}1, 2, 3]> : tensor<3xi32>}> : () -> tensor<?xi32>
// CHECK: return %[[CST]]
%0 = "tfl.transpose"(%cst, %cst_perm) : (tensor<3xi32>, tensor<1xi32>) -> tensor<?xi32>
func.return %0 : tensor<?xi32>
}
// CHECK-LABEL: @add_dense_dense_int_same_shape_dynamic
func.func @add_dense_dense_int_same_shape_dynamic() -> tensor<?xi32> {
%0 = arith.constant dense<[15, 23, -44, -2]> : tensor<4xi32>
%1 = arith.constant dense<[-10, -1, 42, 100]> : tensor<4xi32>
%2 = "tfl.add"(%0, %1) {fused_activation_function = "NONE"} : (tensor<4xi32>, tensor<4xi32>) -> tensor<?xi32>
func.return %2 : tensor<?xi32>
// CHECK: %[[CST:.*]] = "tfl.pseudo_const"() <{value = dense<[5, 22, -2, 98]> : tensor<4xi32>}> : () -> tensor<?xi32>
// CHECK: return %[[CST]]
}
// CHECK-LABEL: @concat_2_tensors_1_empty
func.func @concat_2_tensors_1_empty() -> tensor<2xi32> {
%1 = arith.constant dense<1> : tensor<2xi32>
%2 = arith.constant dense<[]> : tensor<0xi32>
%3 = "tfl.concatenation"(%1, %2) {axis = 0 : i32, fused_activation_function = "NONE"} : (tensor<2xi32>, tensor<0xi32>) -> tensor<2xi32>
func.return %3 : tensor<2xi32>
// CHECK: %[[CST:.*]] = arith.constant dense<1> : tensor<2xi32>
// CHECK: return %[[CST]] : tensor<2xi32>
}
// CHECK-LABEL: @concat_3_tensors_1_empty
func.func @concat_3_tensors_1_empty() -> tensor<?xi32> {
%0 = arith.constant dense<1> : tensor<2xi32>
%1 = arith.constant dense<1> : tensor<2xi32>
%2 = arith.constant dense<[]> : tensor<0xi32>
%3 = "tfl.concatenation"(%0, %1, %2) {axis = 0 : i32, fused_activation_function = "NONE"} : (tensor<2xi32>, tensor<2xi32>, tensor<0xi32>) -> tensor<?xi32>
func.return %3 : tensor<?xi32>
// CHECK: %0 = "tfl.concatenation"(%[[CST]], %[[CST]]) <{axis = 0 : i32, fused_activation_function = "NONE"}>
// CHECK: return %0 : tensor<?xi32>
}
// CHECK-LABEL: @concatConstantTensorsFirstDim
func.func @concatConstantTensorsFirstDim() -> tensor<2x2x3xi32> {
%cst_0 = arith.constant dense<0> : tensor<1x2x3xi32>
%cst_1 = arith.constant dense<1> : tensor<1x2x3xi32>
%0 = "tfl.concatenation"(%cst_0, %cst_1) {axis = 0 : i32, fused_activation_function = "NONE"} : (tensor<1x2x3xi32>, tensor<1x2x3xi32>) -> tensor<2x2x3xi32>
func.return %0 : tensor<2x2x3xi32>
// CHECK: %[[CST:.*]] = arith.constant dense<[{{\[}}{{\[}}0, 0, 0], {{\[}}0, 0, 0]], {{\[}}{{\[}}1, 1, 1], {{\[}}1, 1, 1]]]> : tensor<2x2x3xi32>
// CHECK-NOT: constant-dense
// CHECK-NOT: "tfl.concatenation"
// CHECK: return %[[CST]]
}
// CHECK-LABEL: @concatConstantTensorsMiddleDim
func.func @concatConstantTensorsMiddleDim() -> tensor<1x4x3xi32> {
%cst_0 = arith.constant dense<0> : tensor<1x2x3xi32>
%cst_1 = arith.constant dense<1> : tensor<1x2x3xi32>
%0 = "tfl.concatenation"(%cst_0, %cst_1) {axis = 1 : i32, fused_activation_function = "NONE"} : (tensor<1x2x3xi32>, tensor<1x2x3xi32>) -> tensor<1x4x3xi32>
func.return %0 : tensor<1x4x3xi32>
// CHECK: %[[CST:.*]] = arith.constant dense<[{{\[}}{{\[}}0, 0, 0], {{\[}}0, 0, 0], {{\[}}1, 1, 1], {{\[}}1, 1, 1]]]> : tensor<1x4x3xi32>
// CHECK-NOT: constant-dense
// CHECK-NOT: "tfl.concatenation"
// CHECK: return %[[CST]]
}
// CHECK-LABEL: @concatConstantTensorsLastDim
func.func @concatConstantTensorsLastDim() -> tensor<1x2x6xi32> {
%cst_0 = arith.constant dense<0> : tensor<1x2x3xi32>
%cst_1 = arith.constant dense<1> : tensor<1x2x3xi32>
%0 = "tfl.concatenation"(%cst_0, %cst_1) {axis = 2 : i32, fused_activation_function = "NONE"} : (tensor<1x2x3xi32>, tensor<1x2x3xi32>) -> tensor<1x2x6xi32>
func.return %0 : tensor<1x2x6xi32>
// CHECK: %[[CST:.*]] = arith.constant dense<[{{\[}}{{\[}}0, 0, 0, 1, 1, 1], {{\[}}0, 0, 0, 1, 1, 1]]]> : tensor<1x2x6xi32>
// CHECK-NOT: constant-dense
// CHECK-NOT: "tfl.concatenation"
// CHECK: return %[[CST]]
}
// CHECK-LABEL: @rsqrt_bf16
func.func @rsqrt_bf16() -> tensor<bf16> {
%cst = arith.constant dense<4.0> : tensor<bf16>
%0 = "tfl.rsqrt"(%cst) : (tensor<bf16>) -> tensor<bf16>
func.return %0 : tensor<bf16>
// CHECK: %[[CST:.*]] = arith.constant dense<5.000000e-01> : tensor<bf16>
// CHECK: return %[[CST]]
}
// CHECK-LABEL: @cast_i64_to_i32
func.func @cast_i64_to_i32() -> tensor<5xi32> {
%cst = arith.constant dense<[-1, 0, 1, 2147483647, 2147483648]> : tensor<5xi64>
%0 = "tfl.cast"(%cst) : (tensor<5xi64>) -> tensor<5xi32>
func.return %0 : tensor<5xi32>
// CHECK: %[[CST:.*]] = arith.constant dense<[-1, 0, 1, 2147483647, -2147483648]> : tensor<5xi32>
// CHECK: return %[[CST]]
}
// CHECK-LABEL: @cast_i32_to_i8
func.func @cast_i32_to_i8() -> tensor<6xi8> {
%cst = arith.constant dense<[0, -1, 256, 127, -128, -129]> : tensor<6xi32>
%0 = "tfl.cast"(%cst) : (tensor<6xi32>) -> tensor<6xi8>
func.return %0 : tensor<6xi8>
// CHECK: %[[CST:.*]] = arith.constant dense<[0, -1, 0, 127, -128, 127]> : tensor<6xi8>
// CHECK: return %[[CST]]
}
// CHECK-LABEL: @cast_i8_to_i32
func.func @cast_i8_to_i32() -> tensor<4xi32> {
%cst = arith.constant dense<[0, 128, -1, -128]> : tensor<4xi8>
%0 = "tfl.cast"(%cst) : (tensor<4xi8>) -> tensor<4xi32>
func.return %0 : tensor<4xi32>
// CHECK: %[[CST:.*]] = arith.constant dense<[0, -128, -1, -128]> : tensor<4xi32>
// CHECK: return %[[CST]]
}
// CHECK-LABEL: @cast_identity
func.func @cast_identity(%arg0 : tensor<7xf32>) -> tensor<7xf32> {
%0 = "tfl.cast"(%arg0) : (tensor<7xf32>) -> tensor<7xf32>
func.return %0 : tensor<7xf32>
// CHECK: return %arg0 : tensor<7xf32>
}
// CHECK-LABEL: @cast_i1_to_i8
func.func @cast_i1_to_i8() -> tensor<2xi8> {
%cst = arith.constant dense<[false, true]> : tensor<2xi1>
%0 = "tfl.cast"(%cst) : (tensor<2xi1>) -> tensor<2xi8>
func.return %0 : tensor<2xi8>
// CHECK: %[[CST:.*]] = arith.constant dense<[0, 1]> : tensor<2xi8>
// CHECK: return %[[CST]]
}
// CHECK-LABEL: @cast_i8_to_i1
func.func @cast_i8_to_i1() -> tensor<4xi1> {
%cst = arith.constant dense<[0, 1, 2, -1]> : tensor<4xi8>
%0 = "tfl.cast"(%cst) : (tensor<4xi8>) -> tensor<4xi1>
func.return %0 : tensor<4xi1>
// CHECK: %[[CST:.*]] = arith.constant dense<[false, true, true, true]> : tensor<4xi1>
// CHECK: return %[[CST]]
}
// CHECK-LABEL: @cast_f32_to_i32
func.func @cast_f32_to_i32() -> tensor<8xi32> {
%cst = arith.constant dense<[-1.0, 0.0, 1.5, 0.99, 1.175494351e-38, 3.402823466e+38, -3.402823466e+38, -1.175494351e-38]> : tensor<8xf32>
%0 = "tfl.cast"(%cst) : (tensor<8xf32>) -> tensor<8xi32>
func.return %0 : tensor<8xi32>
}
// CHECK: %cst = arith.constant dense<[-1, 0, 1, 0, 0, 2147483647, -2147483648, 0]> : tensor<8xi32>
// CHECK-LABEL: @cast_f32_to_i64
func.func @cast_f32_to_i64() -> tensor<4xi64> {
%cst = arith.constant dense<[-1.0, 0.0, 1.5, 0.99]> : tensor<4xf32>
%0 = "tfl.cast"(%cst) : (tensor<4xf32>) -> tensor<4xi64>
func.return %0 : tensor<4xi64>
}
// CHECK: %cst = arith.constant dense<[-1, 0, 1, 0]> : tensor<4xi64>
// CHECK-LABEL: @cast_i32_to_f32
func.func @cast_i32_to_f32() -> tensor<5xf32> {
%cst = arith.constant dense<[-1, 0, 2, 2147483647, -2147483648]> : tensor<5xi32>
%0 = "tfl.cast"(%cst) : (tensor<5xi32>) -> tensor<5xf32>
func.return %0 : tensor<5xf32>
}
// CHECK: %cst = arith.constant dense<[-1.000000e+00, 0.000000e+00, 2.000000e+00, 2.14748365E+9, -2.14748365E+9]> : tensor<5xf32>
// CHECK-LABEL: @cast_bool_to_f32
func.func @cast_bool_to_f32() -> tensor<2xf32> {
%cst = arith.constant dense<[true, false]> : tensor<2xi1>
%0 = "tfl.cast"(%cst) : (tensor<2xi1>) -> tensor<2xf32>
func.return %0 : tensor<2xf32>
}
// CHECK: %cst = arith.constant dense<[1.000000e+00, 0.000000e+00]> : tensor<2xf32>
// CHECK-LABEL: @cast_f64_to_f32
func.func @cast_f64_to_f32() -> tensor<4xf32> {
%cst = arith.constant dense<[-1.0, 0.0, 1.5, 100.0]> : tensor<4xf64>
%0 = "tfl.cast"(%cst) : (tensor<4xf64>) -> tensor<4xf32>
func.return %0 : tensor<4xf32>
}
// CHECK: %cst = arith.constant dense<[-1.000000e+00, 0.000000e+00, 1.500000e+00, 1.000000e+02]> : tensor<4xf32>
// CHECK-LABEL: @cast_f32_to_f64
func.func @cast_f32_to_f64() -> tensor<4xf64> {
%cst = arith.constant dense<[-1.0, 0.0, 1.5, 100.0]> : tensor<4xf32>
%0 = "tfl.cast"(%cst) : (tensor<4xf32>) -> tensor<4xf64>
func.return %0 : tensor<4xf64>
}
// CHECK: %cst = arith.constant dense<[-1.000000e+00, 0.000000e+00, 1.500000e+00, 1.000000e+02]> : tensor<4xf64>
// CHECK-LABEL: @cast_f32_to_f16
func.func @cast_f32_to_f16() -> tensor<4xf16> {
%cst = arith.constant dense<[-1.0, 0.0, 1.5, 100.0]> : tensor<4xf32>
%0 = "tfl.cast"(%cst) : (tensor<4xf32>) -> tensor<4xf16>
func.return %0 : tensor<4xf16>
}
// CHECK: %cst = arith.constant dense<[-1.000000e+00, 0.000000e+00, 1.500000e+00, 1.000000e+02]> : tensor<4xf16>
// CHECK-LABEL: @ConstantFoldFullyConnectedSmall
func.func @ConstantFoldFullyConnectedSmall() -> tensor<3xf32> {
%cst_input = arith.constant dense<[2.0, 3.0]> : tensor<2xf32>
%cst_weights = arith.constant dense<[[5.0, 7.0], [11.0, 13.0], [17.0, 19.0]]> : tensor<3x2xf32>
%cst_bias = arith.constant dense<[23.0, 29.0, 31.0]> : tensor<3xf32>
%0 = "tfl.fully_connected" (%cst_input, %cst_weights, %cst_bias) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<2xf32>, tensor<3x2xf32>, tensor<3xf32>) -> tensor<3xf32>
func.return %0 : tensor<3xf32>
// [54, 90, 122]
// CHECK: %[[CST:.*]] = arith.constant dense<[5.400000e+01, 9.000000e+01, 1.220000e+02]> : tensor<3xf32>
// CHECK: return %[[CST]]
}
// CHECK-LABEL: @ConstantFoldFullyConnectedLarge
func.func @ConstantFoldFullyConnectedLarge() -> tensor<1024xf32> {
%cst_input = arith.constant dense<1.0> : tensor<512xf32>
%cst_weights = arith.constant dense<2.0> : tensor<1024x512xf32>
%cst_bias = arith.constant dense<4.0> : tensor<1024xf32>
%0 = "tfl.fully_connected" (%cst_input, %cst_weights, %cst_bias) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<512xf32>, tensor<1024x512xf32>, tensor<1024xf32>) -> tensor<1024xf32>
func.return %0 : tensor<1024xf32>
// 1.0 * 2.0 * 512 + 4.0 = 1028.0
// CHECK: %[[CST:.*]] = arith.constant dense<1.028000e+03> : tensor<1024xf32>
// CHECK: return %[[CST]]
}
// CHECK-LABEL: @ConstantFoldFullyConnectedNoBias
func.func @ConstantFoldFullyConnectedNoBias() -> tensor<1024xf32> {
%cst_input = arith.constant dense<1.0> : tensor<512xf32>
%cst_weights = arith.constant dense<2.0> : tensor<1024x512xf32>
%cst_bias = "tfl.no_value"() {value = unit} : () -> none
%0 = "tfl.fully_connected" (%cst_input, %cst_weights, %cst_bias) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<512xf32>, tensor<1024x512xf32>, none) -> tensor<1024xf32>
func.return %0 : tensor<1024xf32>
// 1.0 * 2.0 * 512 = 1024.0
// CHECK: %[[CST:.*]] = arith.constant dense<1.024000e+03> : tensor<1024xf32>
// CHECK: return %[[CST]]
}
// CHECK-LABEL: @NoFoldFullyConnectedNonFloat
func.func @NoFoldFullyConnectedNonFloat() -> tensor<1024xf32> {
%cst_input = arith.constant dense<1.0> : tensor<512xf32>
%cst_weights = arith.constant dense<2> : tensor<1024x512xi8>
%cst_bias = arith.constant dense<4.0> : tensor<1024xf32>
%0 = "tfl.fully_connected" (%cst_input, %cst_weights, %cst_bias) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<512xf32>, tensor<1024x512xi8>, tensor<1024xf32>) -> tensor<1024xf32>
func.return %0 : tensor<1024xf32>
// CHECK-DAG: %[[CST:.*]] = arith.constant dense<1.000000e+00> : tensor<512xf32>
// CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<2> : tensor<1024x512xi8>
// CHECK-DAG: %[[CST_1:.*]] = arith.constant dense<4.000000e+00> : tensor<1024xf32>
// CHECK: %[[VAL:.*]] = "tfl.fully_connected"(%[[CST]], %[[CST_0]], %[[CST_1]]) <{fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> : (tensor<512xf32>, tensor<1024x512xi8>, tensor<1024xf32>) -> tensor<1024xf32>
// CHECK: return %[[VAL]] : tensor<1024xf32>
}
// CHECK-LABEL: @ConstantFoldFullyConnectedHighRank
func.func @ConstantFoldFullyConnectedHighRank() -> tensor<2x1024xf32> {
%cst_input = arith.constant dense<1.0> : tensor<2x512xf32>
%cst_weights = arith.constant dense<2.0> : tensor<1024x512xf32>
%cst_bias = arith.constant dense<4.0> : tensor<1024xf32>
%0 = "tfl.fully_connected" (%cst_input, %cst_weights, %cst_bias) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<2x512xf32>, tensor<1024x512xf32>, tensor<1024xf32>) -> tensor<2x1024xf32>
func.return %0 : tensor<2x1024xf32>
// 1.0 * 2.0 * 512 + 4.0 = 1028.0
// CHECK: %[[CST:.*]] = arith.constant dense<1.028000e+03> : tensor<2x1024xf32>
// CHECK: return %[[CST]]
}
// CHECK-LABEL: @ConstantFoldFullyConnectedCheckPrecision
func.func @ConstantFoldFullyConnectedCheckPrecision() -> tensor<1xf32> {
%cst_input = arith.constant dense<1.0> : tensor<4xf32>
%cst_weights = arith.constant dense<[[1.0, 1.0e38, 1.0, -1.0e38]]> : tensor<1x4xf32>
%cst_bias = arith.constant dense<0.0> : tensor<1xf32>
%0 = "tfl.fully_connected" (%cst_input, %cst_weights, %cst_bias) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<4xf32>, tensor<1x4xf32>, tensor<1xf32>) -> tensor<1xf32>
func.return %0 : tensor<1xf32>
// CHECK: %[[CST:.*]] = arith.constant dense<2.000000e+00> : tensor<1xf32>
// CHECK: return %[[CST]]
}
// CHECK-LABEL: fully_connected_with_unit_dim
func.func @fully_connected_with_unit_dim() -> tensor<1x5xf32> {
%0 = "tfl.pseudo_const"() <{value = dense<1.0> : tensor<1x5xf32>}> : () -> tensor<1x5xf32>
%1 = "tfl.pseudo_const"() <{value = dense<1.0> : tensor<5x5xf32>}> : () -> tensor<5x5xf32>
%2 = "tfl.pseudo_const"() <{value = dense<1.0> : tensor<1x5xf32>}> : () -> tensor<1x5xf32>
%3 = "tfl.fully_connected"(%0, %1, %2) <{asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = true, weights_format = "DEFAULT"}> : (tensor<1x5xf32>, tensor<5x5xf32>, tensor<1x5xf32>) -> tensor<1x5xf32>
return %3 : tensor<1x5xf32>
}
// CHECK: %cst = arith.constant dense<6.000000e+00> : tensor<1x5xf32>
// CHECK-NOT: fully_connected
// CHECK-LABEL: @ConstantFoldFullyConnectedBatched
func.func @ConstantFoldFullyConnectedBatched() -> tensor<13x1536xf32> {
%cst_input = arith.constant dense<1.0> : tensor<13x1536xf32>
%cst_weights = arith.constant dense<1.0> : tensor<1536x1536xf32>
%cst_bias = "tfl.no_value"() {value = unit} : () -> none
%0 = "tfl.fully_connected" (%cst_input, %cst_weights, %cst_bias) {fused_activation_function = "NONE", keep_num_dims = true, weights_format = "DEFAULT"} : (tensor<13x1536xf32>, tensor<1536x1536xf32>, none) -> tensor<13x1536xf32>
func.return %0 : tensor<13x1536xf32>
// 1.0 * 1.0 * 1536 = 1536.0
// CHECK: %[[CST:.*]] = arith.constant dense<1.536000e+03> : tensor<13x1536xf32>
// CHECK: return %[[CST]]
}
// CHECK-LABEL: @ShapeOpI32
func.func @ShapeOpI32(%arg0 : tensor<576x72xf32>) -> tensor<2xi32> {
%0 = "tfl.shape"(%arg0) : (tensor<576x72xf32>) -> tensor<2xi32>
func.return %0 : tensor<2xi32>
// CHECK: %[[CST:.*]] = arith.constant dense<[576, 72]> : tensor<2xi32>
// CHECK: return %[[CST]]
}
// CHECK-LABEL: @ShapeOpI64
func.func @ShapeOpI64(%arg0 : tensor<576x72xf32>) -> tensor<2xi64> {
%0 = "tfl.shape"(%arg0) : (tensor<576x72xf32>) -> tensor<2xi64>
func.return %0 : tensor<2xi64>
// CHECK: %[[CST:.*]] = arith.constant dense<[576, 72]> : tensor<2xi64>
// CHECK: return %[[CST]]
}
// CHECK-LABEL: @ConstFoldStridedSlice
func.func @ConstFoldStridedSlice(%arg0 : tensor<15600xf32>) -> tensor<15600xf32> {
%0 = "tfl.pseudo_const"() {value = dense<15600> : tensor<1xi32>} : () -> tensor<1xi32>
%1 = "tfl.pseudo_const"() {value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32>
%2 = "tfl.pseudo_const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
%3 = "tfl.strided_slice"(%arg0, %1, %0, %2) {begin_mask = 0 : i32, ellipsis_mask = 0 : i32, end_mask = 0 : i32, new_axis_mask = 0 : i32, shrink_axis_mask = 0 : i32, offset = false} : (tensor<15600xf32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<15600xf32>
func.return %3 : tensor<15600xf32>
// CHECK: return %arg0
}
func.func @ConstFoldStridedSliceMultiDims(%arg0 : tensor<10x10x10xf32>) -> tensor<10x10x10xf32> {
%0 = "tfl.pseudo_const"() {value = dense<[10, 10, 10]> : tensor<3xi32>} : () -> tensor<3xi32>
%1 = "tfl.pseudo_const"() {value = dense<0> : tensor<3xi32>} : () -> tensor<3xi32>
%2 = "tfl.pseudo_const"() {value = dense<1> : tensor<3xi32>} : () -> tensor<3xi32>
%3 = "tfl.strided_slice"(%arg0, %1, %0, %2) {begin_mask = 0 : i32, ellipsis_mask = 0 : i32, end_mask = 0 : i32, new_axis_mask = 0 : i32, shrink_axis_mask = 0 : i32, offset = false} : (tensor<10x10x10xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<10x10x10xf32>
func.return %3 : tensor<10x10x10xf32>
// CHECK: return %arg0
}
func.func @NotFoldStridedSlice(%arg0 : tensor<10x10x10xf32>) -> tensor<9x9x9xf32> {
%0 = "tfl.pseudo_const"() {value = dense<[9, 9, 9]> : tensor<3xi32>} : () -> tensor<3xi32>
%1 = "tfl.pseudo_const"() {value = dense<0> : tensor<3xi32>} : () -> tensor<3xi32>
%2 = "tfl.pseudo_const"() {value = dense<1> : tensor<3xi32>} : () -> tensor<3xi32>
%3 = "tfl.strided_slice"(%arg0, %1, %0, %2) {begin_mask = 0 : i32, ellipsis_mask = 0 : i32, end_mask = 0 : i32, new_axis_mask = 0 : i32, shrink_axis_mask = 0 : i32, offset = false} : (tensor<10x10x10xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<9x9x9xf32>
func.return %3 : tensor<9x9x9xf32>
// CHECK: %[[STRIDED_SLICE:.*]] = "tfl.strided_slice"
// CHECK: return %[[STRIDED_SLICE]]
}
func.func @ConstFoldPad(%arg0: tensor<15600xf32>) -> tensor<15600xf32> {
%0 = "tfl.pseudo_const"() {value = dense<0> : tensor<1x2xi32>} : () -> tensor<1x2xi32>
%1 = "tfl.pad"(%arg0, %0) : (tensor<15600xf32>, tensor<1x2xi32>) -> tensor<15600xf32>
func.return %1 : tensor<15600xf32>
// CHECK: return %arg0
}
func.func @ConstFoldPadV2(%arg0: tensor<15600xf32>) -> tensor<15600xf32> {
%0 = "tfl.pseudo_const"() {value = dense<0> : tensor<1x2xi32>} : () -> tensor<1x2xi32>
%1 = "tfl.pseudo_const"() {value = dense<0.0> : tensor<f32>} : () -> tensor<f32>
%2 = "tfl.padv2"(%arg0, %0, %1) : (tensor<15600xf32>, tensor<1x2xi32>, tensor<f32>) -> tensor<15600xf32>
func.return %2 : tensor<15600xf32>
// CHECK: return %arg0
}
// CHECK-LABEL: @ConstFoldEmbeddingLookup
func.func @ConstFoldEmbeddingLookup() -> (tensor<5x2xf32>, tensor<3x2x2xf32>) {
%index0 = "tfl.pseudo_const"() {value = dense<[2, 1, 0, 0, 2]> : tensor<5xi32>} : () -> tensor<5xi32>
%index1 = "tfl.pseudo_const"() {value = dense<[0, 1, 0]> : tensor<3xi32>} : () -> tensor<3xi32>
%value0 = "tfl.pseudo_const"() {value = dense<[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]> : tensor<3x2xf32>} : () -> tensor<3x2xf32>
%value1 = "tfl.pseudo_const"() {value = dense<[[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]]> : tensor<2x2x2xf32>} : () -> tensor<2x2x2xf32>
%lookup0 = "tfl.embedding_lookup"(%index0, %value0) : (tensor<5xi32>, tensor<3x2xf32>) -> tensor<5x2xf32>
%lookup1 = "tfl.embedding_lookup"(%index1, %value1) : (tensor<3xi32>, tensor<2x2x2xf32>) -> tensor<3x2x2xf32>
func.return %lookup0, %lookup1 : tensor<5x2xf32>, tensor<3x2x2xf32>
// CHECK-DAG: %[[LOOKUP0:.*]] = arith.constant dense<{{\[\[}}5.000000e+00, 6.000000e+00], [3.000000e+00, 4.000000e+00], [1.000000e+00, 2.000000e+00], [1.000000e+00, 2.000000e+00], [5.000000e+00, 6.000000e+00]]> : tensor<5x2xf32>
// CHECK-DAG: %[[LOOKUP1:.*]] = arith.constant dense<{{\[\[\[}}1.000000e+00, 2.000000e+00], [3.000000e+00, 4.000000e+00]], {{\[\[}}5.000000e+00, 6.000000e+00], [7.000000e+00, 8.000000e+00]], {{\[\[}}1.000000e+00, 2.000000e+00], [3.000000e+00, 4.000000e+00]]]> : tensor<3x2x2xf32>
// CHECK: return %[[LOOKUP0]], %[[LOOKUP1]] : tensor<5x2xf32>, tensor<3x2x2xf32>
}
// CHECK-LABEL: @less_int_both_splat
func.func @less_int_both_splat() -> tensor<4xi1> {
%0 = arith.constant dense<3> : tensor<4xi32>
%1 = arith.constant dense<10> : tensor<4xi32>
%2 = "tfl.less"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
func.return %2 : tensor<4xi1>
}
// CHECK: %cst = arith.constant dense<true> : tensor<4xi1>
// CHECK-LABEL: @less_int_one_splat
func.func @less_int_one_splat() -> tensor<4xi1> {
%0 = arith.constant dense<3> : tensor<4xi32>
%1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi32>
%2 = "tfl.less"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
func.return %2 : tensor<4xi1>
}
// CHECK:%cst = arith.constant dense<[true, false, false, false]> : tensor<4xi1>
// CHECK-LABEL: @less_int
func.func @less_int() -> tensor<4xi1> {
%0 = arith.constant dense<[11, 2, 0, 2]> : tensor<4xi32>
%1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi32>
%2 = "tfl.less"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
func.return %2 : tensor<4xi1>
}
// CHECK: %cst = arith.constant dense<[false, false, false, true]> : tensor<4xi1>
// CHECK-LABEL: @less_int64
func.func @less_int64() -> tensor<4xi1> {
%0 = arith.constant dense<[11, 2, 0, 2]> : tensor<4xi64>
%1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi64>
%2 = "tfl.less"(%0, %1) : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi1>
func.return %2 : tensor<4xi1>
}
// CHECK: %cst = arith.constant dense<[false, false, false, true]> : tensor<4xi1>
// CHECK-LABEL: @less_float
func.func @less_float() -> tensor<4xi1> {
%0 = arith.constant dense<[11.0, 2.0, 0.0, 2.0]> : tensor<4xf32>
%1 = arith.constant dense<[10.0, 2.0, -1.0, 3.0]> : tensor<4xf32>
%2 = "tfl.less"(%0, %1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xi1>
func.return %2 : tensor<4xi1>
}
// CHECK: %cst = arith.constant dense<[false, false, false, true]> : tensor<4xi1>
// CHECK-LABEL: @less_equal_int
func.func @less_equal_int() -> tensor<4xi1> {
%0 = arith.constant dense<[11, 2, 0, 2]> : tensor<4xi32>
%1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi32>
%2 = "tfl.less_equal"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
func.return %2 : tensor<4xi1>
}
// CHECK: %cst = arith.constant dense<[false, true, false, true]> : tensor<4xi1>
// CHECK-LABEL: @less_equal_int64
func.func @less_equal_int64() -> tensor<4xi1> {
%0 = arith.constant dense<[11, 2, 0, 2]> : tensor<4xi64>
%1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi64>
%2 = "tfl.less_equal"(%0, %1) : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi1>
func.return %2 : tensor<4xi1>
}
// CHECK: %cst = arith.constant dense<[false, true, false, true]> : tensor<4xi1>
// CHECK-LABEL: @less_equal_float
func.func @less_equal_float() -> tensor<4xi1> {
%0 = arith.constant dense<[11.0, 2.0, 0.0, 2.0]> : tensor<4xf32>
%1 = arith.constant dense<[10.0, 2.0, -1.0, 3.0]> : tensor<4xf32>
%2 = "tfl.less_equal"(%0, %1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xi1>
func.return %2 : tensor<4xi1>
}
// CHECK: %cst = arith.constant dense<[false, true, false, true]> : tensor<4xi1>
// CHECK-LABEL: @greater_int
func.func @greater_int() -> tensor<4xi1> {
%0 = arith.constant dense<[11, 2, 0, 2]> : tensor<4xi32>
%1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi32>
%2 = "tfl.greater"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
func.return %2 : tensor<4xi1>
}
// CHECK: %cst = arith.constant dense<[true, false, true, false]> : tensor<4xi1>
// CHECK-LABEL: @greater_int64
func.func @greater_int64() -> tensor<4xi1> {
%0 = arith.constant dense<[11, 2, 0, 2]> : tensor<4xi64>
%1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi64>
%2 = "tfl.greater"(%0, %1) : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi1>
func.return %2 : tensor<4xi1>
}
// CHECK: %cst = arith.constant dense<[true, false, true, false]> : tensor<4xi1>
// CHECK-LABEL: @greater_float
func.func @greater_float() -> tensor<4xi1> {
%0 = arith.constant dense<[11.0, 2.0, 0.0, 2.0]> : tensor<4xf32>
%1 = arith.constant dense<[10.0, 2.0, -1.0, 3.0]> : tensor<4xf32>
%2 = "tfl.greater"(%0, %1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xi1>
func.return %2 : tensor<4xi1>
}
// CHECK: %cst = arith.constant dense<[true, false, true, false]> : tensor<4xi1>
// CHECK-LABEL: @greater_equal_int
func.func @greater_equal_int() -> tensor<4xi1> {
%0 = arith.constant dense<[11, 2, 0, 2]> : tensor<4xi32>
%1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi32>
%2 = "tfl.greater_equal"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
func.return %2 : tensor<4xi1>
}
// CHECK: %cst = arith.constant dense<[true, true, true, false]> : tensor<4xi1>
// CHECK-LABEL: @greater_equal_int64
func.func @greater_equal_int64() -> tensor<4xi1> {
%0 = arith.constant dense<[11, 2, 0, 2]> : tensor<4xi64>
%1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi64>
%2 = "tfl.greater_equal"(%0, %1) : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi1>
func.return %2 : tensor<4xi1>
}
// CHECK: %cst = arith.constant dense<[true, true, true, false]> : tensor<4xi1>
// CHECK-LABEL: @greater_equal_float
func.func @greater_equal_float() -> tensor<4xi1> {
%0 = arith.constant dense<[11.0, 2.0, 0.0, 2.0]> : tensor<4xf32>
%1 = arith.constant dense<[10.0, 2.0, -1.0, 3.0]> : tensor<4xf32>
%2 = "tfl.greater_equal"(%0, %1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xi1>
func.return %2 : tensor<4xi1>
}
// CHECK: %cst = arith.constant dense<[true, true, true, false]> : tensor<4xi1>
// CHECK-LABEL: @equal_int
func.func @equal_int() -> tensor<4xi1> {
%0 = arith.constant dense<[11, 2, 0, 2]> : tensor<4xi32>
%1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi32>
%2 = "tfl.equal"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
func.return %2 : tensor<4xi1>
}
// CHECK: %cst = arith.constant dense<[false, true, false, false]> : tensor<4xi1>
// CHECK-LABEL: @equal_int64
func.func @equal_int64() -> tensor<4xi1> {
%0 = arith.constant dense<[11, 2, 0, 2]> : tensor<4xi64>
%1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi64>
%2 = "tfl.equal"(%0, %1) : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi1>
func.return %2 : tensor<4xi1>
}
// CHECK: %cst = arith.constant dense<[false, true, false, false]> : tensor<4xi1>
// CHECK-LABEL: @equal_float
func.func @equal_float() -> tensor<4xi1> {
%0 = arith.constant dense<[11.0, 2.0, 0.0, 2.0]> : tensor<4xf32>
%1 = arith.constant dense<[10.0, 2.0, -1.0, 3.0]> : tensor<4xf32>
%2 = "tfl.equal"(%0, %1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xi1>
func.return %2 : tensor<4xi1>
}
// CHECK: %cst = arith.constant dense<[false, true, false, false]> : tensor<4xi1>
// CHECK-LABEL: @not_equal_int
func.func @not_equal_int() -> tensor<4xi1> {
%0 = arith.constant dense<[11, 2, 0, 2]> : tensor<4xi32>
%1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi32>
%2 = "tfl.not_equal"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
func.return %2 : tensor<4xi1>
}
// CHECK: %cst = arith.constant dense<[true, false, true, true]> : tensor<4xi1>
// CHECK-LABEL: @not_equal_int64
func.func @not_equal_int64() -> tensor<4xi1> {
%0 = arith.constant dense<[11, 2, 0, 2]> : tensor<4xi64>
%1 = arith.constant dense<[10, 2, -1, 3]> : tensor<4xi64>
%2 = "tfl.not_equal"(%0, %1) : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi1>
func.return %2 : tensor<4xi1>
}
// CHECK: %cst = arith.constant dense<[true, false, true, true]> : tensor<4xi1>
// CHECK-LABEL: @not_equal_float
func.func @not_equal_float() -> tensor<4xi1> {
%0 = arith.constant dense<[11.0, 2.0, 0.0, 2.0]> : tensor<4xf32>
%1 = arith.constant dense<[10.0, 2.0, -1.0, 3.0]> : tensor<4xf32>
%2 = "tfl.not_equal"(%0, %1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xi1>
func.return %2 : tensor<4xi1>
}
// CHECK: %cst = arith.constant dense<[true, false, true, true]> : tensor<4xi1>
// CHECK-LABEL: @logical_or
func.func @logical_or() -> tensor<3xi1> {
%0 = arith.constant dense<[true, false, true]> : tensor<3xi1>
%1 = arith.constant dense<[false, false, true]> : tensor<3xi1>
%2 = "tfl.logical_or"(%0, %1) : (tensor<3xi1>, tensor<3xi1>) -> tensor<3xi1>
func.return %2 : tensor<3xi1>
}
// CHECK: %cst = arith.constant dense<[true, false, true]> : tensor<3xi1>
// CHECK-LABEL: @logical_and
func.func @logical_and() -> tensor<3xi1> {
%0 = arith.constant dense<[true, false, true]> : tensor<3xi1>
%1 = arith.constant dense<[false, false, true]> : tensor<3xi1>
%2 = "tfl.logical_and"(%0, %1) : (tensor<3xi1>, tensor<3xi1>) -> tensor<3xi1>
func.return %2 : tensor<3xi1>
}
// CHECK: %cst = arith.constant dense<[false, false, true]> : tensor<3xi1>
// CHECK-LABEL: @select_splat_cond
func.func @select_splat_cond() -> tensor<4xi32> {
%cond = arith.constant dense<true> : tensor<4xi1>
%0 = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32>
%1 = arith.constant dense<[-1, -2, -3, -4]> : tensor<4xi32>
%2 = "tfl.select"(%cond, %0, %1) : (tensor<4xi1>, tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32>
func.return %2 : tensor<4xi32>
}
// CHECK: %cst = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32>
// CHECK-LABEL: select_splat_lhs
func.func @select_splat_lhs() -> tensor<4xi32> {
%cond = arith.constant dense<[true, true, false, false]> : tensor<4xi1>
%0 = arith.constant dense<0> : tensor<4xi32>
%1 = arith.constant dense<[-1, -2, -3, -4]> : tensor<4xi32>
%2 = "tfl.select"(%cond, %0, %1) : (tensor<4xi1>, tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32>
func.return %2 : tensor<4xi32>
}
// CHECK: %cst = arith.constant dense<[0, 0, -3, -4]> : tensor<4xi32>
// CHECK-LABEL: select_float
func.func @select_float() -> tensor<4xf32> {
%cond = arith.constant dense<[true, true, false, false]> : tensor<4xi1>
%0 = arith.constant dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32>
%1 = arith.constant dense<[-1.0, -2.0, -3.0, -4.0]> : tensor<4xf32>
%2 = "tfl.select"(%cond, %0, %1) : (tensor<4xi1>, tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
func.return %2 : tensor<4xf32>
}
// CHECK: %cst = arith.constant dense<[1.000000e+00, 2.000000e+00, -3.000000e+00, -4.000000e+00]> : tensor<4xf32
// CHECK-LABEL: ceil
func.func @ceil() -> tensor<3xf32> {
%cst = arith.constant dense<[-1.0, 0.0, 0.99]> : tensor<3xf32>
%0 = "tfl.ceil"(%cst) : (tensor<3xf32>) -> tensor<3xf32>
func.return %0 : tensor<3xf32>
}
// CHECK: %cst = arith.constant dense<[-1.000000e+00, 0.000000e+00, 1.000000e+00]> : tensor<3xf32>
// CHECK-LABEL: ceil_f64
func.func @ceil_f64() -> tensor<3xf64> {
%cst = arith.constant dense<[-1.0, 0.0, 0.99]> : tensor<3xf64>
%0 = "tfl.ceil"(%cst) : (tensor<3xf64>) -> tensor<3xf64>
func.return %0 : tensor<3xf64>
}
// CHECK: tfl.ceil
// CHECK-LABEL: floor
func.func @floor() -> tensor<3xf32> {
%cst = arith.constant dense<[-1.0, 0.0, 0.99]> : tensor<3xf32>
%0 = "tfl.floor"(%cst) : (tensor<3xf32>) -> tensor<3xf32>
func.return %0 : tensor<3xf32>
}
// CHECK: %cst = arith.constant dense<[-1.000000e+00, 0.000000e+00, 0.000000e+00]> : tensor<3xf32>
// CHECK-LABEL: floor_f64
func.func @floor_f64() -> tensor<3xf64> {
%cst = arith.constant dense<[-1.0, 0.0, 0.99]> : tensor<3xf64>
%0 = "tfl.floor"(%cst) : (tensor<3xf64>) -> tensor<3xf64>
func.return %0 : tensor<3xf64>
}
// CHECK: tfl.floor
// CHECK-LABEL: exp
func.func @exp() -> tensor<4xf32> {
%cst = arith.constant dense<[-1.0, 0.0, 0.99, 0.36787944117]> : tensor<4xf32>
%0 = "tfl.exp"(%cst) : (tensor<4xf32>) -> tensor<4xf32>
func.return %0 : tensor<4xf32>
}
// CHECK: %cst = arith.constant dense<[0.36787945, 1.000000e+00, 2.69123459, 1.44466782]> : tensor<4xf32>
// CHECK-LABEL: exp_f64
func.func @exp_f64() -> tensor<4xf64> {
%cst = arith.constant dense<[-1.0, 0.0, 0.99, 0.36787944117]> : tensor<4xf64>
%0 = "tfl.exp"(%cst) : (tensor<4xf64>) -> tensor<4xf64>
func.return %0 : tensor<4xf64>
}
// CHECK: tfl.exp
// CHECK-LABEL: pow_float
func.func @pow_float() -> tensor<3xf32> {
%0 = arith.constant dense<[1.0, 0.0, 2.0]> : tensor<3xf32>
%1 = arith.constant dense<[2.0, 3.0, -1.5]> : tensor<3xf32>
%2 = "tfl.pow"(%0, %1) : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32>
func.return %2 : tensor<3xf32>
}
// CHECK: %cst = arith.constant dense<[1.000000e+00, 0.000000e+00, 0.353553385]> : tensor<3xf32>
// CHECK-LABEL: pow_int
func.func @pow_int() -> tensor<3xi32> {
%0 = arith.constant dense<[1, 0, 2]> : tensor<3xi32>
%1 = arith.constant dense<[2, 3, -1]> : tensor<3xi32>
%2 = "tfl.pow"(%0, %1) : (tensor<3xi32>, tensor<3xi32>) -> tensor<3xi32>
func.return %2 : tensor<3xi32>
}
// CHECK: %cst = arith.constant dense<[1, 0, 0]> : tensor<3xi32>
// CHECK-LABEL: logical_not
func.func @logical_not() -> tensor<3xi1> {
%cst = arith.constant dense<[false, true, false]> : tensor<3xi1>
%0 = "tfl.logical_not"(%cst) : (tensor<3xi1>) -> tensor<3xi1>
func.return %0 : tensor<3xi1>
}
// CHECK: %cst = arith.constant dense<[true, false, true]> : tensor<3xi1>
// CHECK-LABEL: logical_not_splat
func.func @logical_not_splat() -> tensor<3xi1> {
%cst = arith.constant dense<false> : tensor<3xi1>
%0 = "tfl.logical_not"(%cst) : (tensor<3xi1>) -> tensor<3xi1>
func.return %0 : tensor<3xi1>
}
// CHECK: %cst = arith.constant dense<true> : tensor<3xi1>
// CHECK-LABEL: bitwise_xor_i32
func.func @bitwise_xor_i32() -> tensor<3xi32> {
%0 = arith.constant dense<[0, 5, 3]> : tensor<3xi32>
%1 = arith.constant dense<[5, 0, 7]> : tensor<3xi32>
%2 = "tfl.bitwise_xor"(%0, %1) : (tensor<3xi32>, tensor<3xi32>) -> tensor<3xi32>
func.return %2 : tensor<3xi32>
}
// CHECK: %cst = arith.constant dense<[5, 5, 4]> : tensor<3xi32>
// CHECK-LABEL: bitwise_xor_i8
func.func @bitwise_xor_i8() -> tensor<3xi8> {
%0 = arith.constant dense<[0, 5, 3]> : tensor<3xi8>
%1 = arith.constant dense<[5, 0, 7]> : tensor<3xi8>
%2 = "tfl.bitwise_xor"(%0, %1) : (tensor<3xi8>, tensor<3xi8>) -> tensor<3xi8>
func.return %2 : tensor<3xi8>
}
// CHECK: %cst = arith.constant dense<[5, 5, 4]> : tensor<3xi8>
// CHECK-LABEL: relu
func.func @relu() -> tensor<3xf32> {
%cst = arith.constant dense<[-1.0, 0.0, 0.99]> : tensor<3xf32>
%0 = "tfl.relu"(%cst) : (tensor<3xf32>) -> tensor<3xf32>
func.return %0 : tensor<3xf32>
}
// CHECK: %cst = arith.constant dense<[0.000000e+00, 0.000000e+00, 9.900000e-01]> : tensor<3xf32>
// CHECK-LABEL: slice
func.func @slice_first_dim() -> tensor<1x1x5x6xf32> {
%cst_0 = arith.constant dense<9.000000e+00> : tensor<2x1x5x6xf32>
%cst_1 = arith.constant dense<0> : tensor<4xi32>
%cst_2 = arith.constant dense<[1, 1, 5, 6]> : tensor<4xi32>
%0 = "tfl.slice"(%cst_0, %cst_1, %cst_2) : (tensor<2x1x5x6xf32>, tensor<4xi32>, tensor<4xi32>) -> tensor<1x1x5x6xf32>
func.return %0 : tensor<1x1x5x6xf32>
}
// CHECK %cst = arith.constant dense<9.000000e+00> : tensor<1x1x5x6xf32>
// CHECK-LABEL: slice_trivial
func.func @slice_trivial(%arg0: tensor<2x1x5x6xf32>) -> tensor<2x1x5x6xf32> {
%cst_1 = arith.constant dense<0> : tensor<4xi32>
%cst_2 = arith.constant dense<[2, 1, 5, 6]> : tensor<4xi32>
%0 = "tfl.slice"(%arg0, %cst_1, %cst_2) : (tensor<2x1x5x6xf32>, tensor<4xi32>, tensor<4xi32>) -> tensor<2x1x5x6xf32>
func.return %0 : tensor<2x1x5x6xf32>
}
// CHECK-NOT: tfl.slice
// CHECK-LABEL: sum
func.func @sum() -> tensor<2xf32> {
%cst = arith.constant dense<[0, 1]> : tensor<2xi32>
%cst_1 = arith.constant dense<[[[0.0, 1.0], [2.0, 3.0]], [[4.0, 5.0], [6.0, 7.0]]]> : tensor<2x2x2xf32>
%0 = "tfl.sum"(%cst_1, %cst) <{keep_dims = false}> : (tensor<2x2x2xf32>, tensor<2xi32>) -> tensor<2xf32>
func.return %0 : tensor<2xf32>
}
// CHECK: arith.constant dense<[1.200000e+01, 1.600000e+01]> : tensor<2xf32>
// CHECK-LABEL: sum_keep_dims
func.func @sum_keep_dims() -> tensor<1x1x2xf32> {
%cst = arith.constant dense<[0, 1]> : tensor<2xi32>
%cst_1 = arith.constant dense<[[[0.0, 1.0], [2.0, 3.0]], [[4.0, 5.0], [6.0, 7.0]]]> : tensor<2x2x2xf32>
%0 = "tfl.sum"(%cst_1, %cst) <{keep_dims = true}> : (tensor<2x2x2xf32>, tensor<2xi32>) -> tensor<1x1x2xf32>
func.return %0 : tensor<1x1x2xf32>
}
// CHECK-LITERAL: arith.constant dense<[[[1.200000e+01, 1.600000e+01]]]> : tensor<1x1x2xf32>
// CHECK-LABEL: sum_general_reduction_dims
func.func @sum_general_reduction_dims() -> tensor<3xf32> {
%cst = arith.constant dense<[0, 2]> : tensor<2xi32>
%cst_1 = arith.constant dense<2.000000e+00> : tensor<1x3x2xf32>
%0 = "tfl.sum"(%cst_1, %cst) <{keep_dims = false}> : (tensor<1x3x2xf32>, tensor<2xi32>) -> tensor<3xf32>
func.return %0 : tensor<3xf32>
}
// CHECK: %cst = arith.constant dense<4.000000e+00> : tensor<3xf32>
// CHECK-LABEL: sum_general_reduction_dims_keep_dims
func.func @sum_general_reduction_dims_keep_dims() -> tensor<1x3x1xf32> {
%cst = arith.constant dense<[0, 2]> : tensor<2xi32>
%cst_1 = arith.constant dense<2.000000e+00> : tensor<1x3x2xf32>
%0 = "tfl.sum"(%cst_1, %cst) <{keep_dims = true}> : (tensor<1x3x2xf32>, tensor<2xi32>) -> tensor<1x3x1xf32>
func.return %0 : tensor<1x3x1xf32>
}
// CHECK: %cst = arith.constant dense<4.000000e+00> : tensor<1x3x1xf32>
// CHECK-LABEL: gather
func.func @gather() -> (tensor<2x3x4x5xi16>, tensor<2x3x4x5xi16>) {
%params = arith.constant dense<[
[[[1111, 1112, 1113, 1114, 1115],
[1121, 1122, 1123, 1124, 1125],
[1131, 1132, 1133, 1134, 1135],
[1141, 1142, 1143, 1144, 1145],
[1151, 1152, 1153, 1154, 1155],
[1161, 1162, 1163, 1164, 1165]],
[[1211, 1212, 1213, 1214, 1215],
[1221, 1222, 1223, 1224, 1225],
[1231, 1232, 1233, 1234, 1235],
[1241, 1242, 1243, 1244, 1245],
[1251, 1252, 1253, 1254, 1255],
[1261, 1262, 1263, 1264, 1265]],
[[1311, 1312, 1313, 1314, 1315],
[1321, 1322, 1323, 1324, 1325],
[1331, 1332, 1333, 1334, 1335],
[1341, 1342, 1343, 1344, 1345],
[1351, 1352, 1353, 1354, 1355],
[1361, 1362, 1363, 1364, 1365]]],
[[[2111, 2112, 2113, 2114, 2115],
[2121, 2122, 2123, 2124, 2125],
[2131, 2132, 2133, 2134, 2135],
[2141, 2142, 2143, 2144, 2145],
[2151, 2152, 2153, 2154, 2155],
[2161, 2162, 2163, 2164, 2165]],
[[2211, 2212, 2213, 2214, 2215],
[2221, 2222, 2223, 2224, 2225],
[2231, 2232, 2233, 2234, 2235],
[2241, 2242, 2243, 2244, 2245],
[2251, 2252, 2253, 2254, 2255],
[2261, 2262, 2263, 2264, 2265]],
[[2311, 2312, 2313, 2314, 2315],
[2321, 2322, 2323, 2324, 2325],
[2331, 2332, 2333, 2334, 2335],
[2341, 2342, 2343, 2344, 2345],
[2351, 2352, 2353, 2354, 2355],
[2361, 2362, 2363, 2364, 2365]]]]> : tensor<2x3x6x5xi16>
%indices = arith.constant dense<[[5, 4, 3, 2], [3, 2, 1, 0]]> : tensor<2x4xi64>
%gathered = "tfl.gather"(%params, %indices) <{axis = 2 : i32, batch_dims = 1 : i32}> : (tensor<2x3x6x5xi16>, tensor<2x4xi64>) -> tensor<2x3x4x5xi16>
// This is the same tensor as the one on the CHECK line
%expected = arith.constant dense<[
[[[1161, 1162, 1163, 1164, 1165],
[1151, 1152, 1153, 1154, 1155],
[1141, 1142, 1143, 1144, 1145],
[1131, 1132, 1133, 1134, 1135]],
[[1261, 1262, 1263, 1264, 1265],
[1251, 1252, 1253, 1254, 1255],
[1241, 1242, 1243, 1244, 1245],
[1231, 1232, 1233, 1234, 1235]],
[[1361, 1362, 1363, 1364, 1365],
[1351, 1352, 1353, 1354, 1355],
[1341, 1342, 1343, 1344, 1345],
[1331, 1332, 1333, 1334, 1335]]],
[[[2141, 2142, 2143, 2144, 2145],
[2131, 2132, 2133, 2134, 2135],
[2121, 2122, 2123, 2124, 2125],
[2111, 2112, 2113, 2114, 2115]],
[[2241, 2242, 2243, 2244, 2245],
[2231, 2232, 2233, 2234, 2235],
[2221, 2222, 2223, 2224, 2225],
[2211, 2212, 2213, 2214, 2215]],
[[2341, 2342, 2343, 2344, 2345],
[2331, 2332, 2333, 2334, 2335],
[2321, 2322, 2323, 2324, 2325],
[2311, 2312, 2313, 2314, 2315]]]]> : tensor<2x3x4x5xi16>
func.return %gathered, %expected : tensor<2x3x4x5xi16>, tensor<2x3x4x5xi16>
// CHECK-NOT: tfl.gather
// CHECK: [[CST:%.*]] = arith.constant dense<"0x89048A048B048C048D047F048004810482048304750476047704780479046B046C046D046E046F04ED04EE04EF04F004F104E304E404E504E604E704D904DA04DB04DC04DD04CF04D004D104D204D304510552055305540555054705480549054A054B053D053E053F0540054105330534053505360537055D085E085F08600861085308540855085608570849084A084B084C084D083F084008410842084308C108C208C308C408C508B708B808B908BA08BB08AD08AE08AF08B008B108A308A408A508A608A708250926092709280929091B091C091D091E091F09110912091309140915090709080909090A090B09"> : tensor<2x3x4x5xi16>
// If the return value is the same constant twice, the result is the same as expected
// CHECK: return [[CST]], [[CST]]
}
// CHECK-LABEL: func @gather_nd_slices
func.func @gather_nd_slices() -> tensor<2x2xi32> {
%params = arith.constant dense<[[[1, 2], [3, 4]], [[5, 6], [7, 8]]]> : tensor<2x2x2xi32>
%indices = arith.constant dense<[[0, 1], [1, 0]]> : tensor<2x2xi64>
%0 = "tfl.gather_nd"(%params, %indices) : (tensor<2x2x2xi32>, tensor<2x2xi64>) -> tensor<2x2xi32>
return %0 : tensor<2x2xi32>
// CHECK-NOT: tfl.gather_nd
// CHECK: [[CST:%.*]] = arith.constant dense<{{\[\[}}3, 4], {{\[}}5, 6]]> : tensor<2x2xi32>
// CHECK: return [[CST]]
}
// CHECK-LABEL: func @gather_nd_scalars
func.func @gather_nd_scalars() -> tensor<4xf32> {
%params = arith.constant dense<[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]> : tensor<3x3xf32>
%indices = arith.constant dense<[[0, 0], [2, 2], [1, 0], [2, 0]]> : tensor<4x2xi32>
%0 = "tfl.gather_nd"(%params, %indices) : (tensor<3x3xf32>, tensor<4x2xi32>) -> tensor<4xf32>
return %0 : tensor<4xf32>
// CHECK-NOT: tfl.gather_nd
// CHECK: [[CST:%.+]] = arith.constant dense<[1.000000e+00, 9.000000e+00, 4.000000e+00, 7.000000e+00]> : tensor<4xf32>
// CHECK: return [[CST]]
}
// CHECK-LABEL: reverse_2_dims
func.func @reverse_2_dims() -> tensor<2x3x2xi32> {
%input = "tfl.pseudo_const"() <{value = dense<[[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12]]]> : tensor<2x3x2xi32>}> : () -> tensor<2x3x2xi32>
%axis = "tfl.pseudo_const"() <{value = dense<[0, 1]> : tensor<2xi32>}> : () -> tensor<2xi32>
%reverse = "tfl.reverse_v2"(%input, %axis) : (tensor<2x3x2xi32>, tensor<2xi32>) -> tensor<2x3x2xi32>
return %reverse : tensor<2x3x2xi32>
}
// CHECK-LITERAL: %cst = arith.constant dense<[[[11, 12], [9, 10], [7, 8]], [[5, 6], [3, 4], [1, 2]]]> : tensor<2x3x2xi32>
// CHECK: return %cst : tensor<2x3x2xi32>
// CHECK-LABEL: reverse_1_dim
func.func @reverse_1_dim() -> tensor<2x3x2xi32> {
%input = "tfl.pseudo_const"() <{value = dense<[[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12]]]> : tensor<2x3x2xi32>}> : () -> tensor<2x3x2xi32>
%axis = "tfl.pseudo_const"() <{value = dense<[2]> : tensor<1xi32>}> : () -> tensor<1xi32>
%reverse = "tfl.reverse_v2"(%input, %axis) : (tensor<2x3x2xi32>, tensor<1xi32>) -> tensor<2x3x2xi32>
return %reverse : tensor<2x3x2xi32>
}
// CHECK-LITERAL: %cst = arith.constant dense<[[[2, 1], [4, 3], [6, 5]], [[8, 7], [10, 9], [12, 11]]]> : tensor<2x3x2xi32>
// CHECK: return %cst : tensor<2x3x2xi32>
// CHECK-LABEL: @slice_no_op
func.func @slice_no_op() -> tensor<2x3x2xi32> {
%cst_1 = arith.constant dense<[[[0, 1], [2, 3], [4, 5]], [[6, 7], [8, 9], [10, 11]]]> : tensor<2x3x2xi32>
%cst_2 = arith.constant dense<0> : tensor<3xi32>
%cst_3 = arith.constant dense<[2, 3, 2]> : tensor<3xi32>
%0 = "tfl.slice"(%cst_1, %cst_2, %cst_3) : (tensor<2x3x2xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<2x3x2xi32>
return %0 : tensor<2x3x2xi32>
}
// CHECK-LITERAL: %cst = arith.constant dense<[[[0, 1], [2, 3], [4, 5]], [[6, 7], [8, 9], [10, 11]]]> : tensor<2x3x2xi32>
// CHECK-NOT: slice
// CHECK-LABEL: @slice_some_dims
func.func @slice_some_dims() -> tensor<2x2x1xi32> {
%cst_1 = arith.constant dense<[[[0, 1], [2, 3], [4, 5]], [[6, 7], [8, 9], [10, 11]]]> : tensor<2x3x2xi32>
%cst_2 = arith.constant dense<[0, 1, 1]> : tensor<3xi32>
%cst_3 = arith.constant dense<[2, 2, 1]> : tensor<3xi32>
%0 = "tfl.slice"(%cst_1, %cst_2, %cst_3) : (tensor<2x3x2xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<2x2x1xi32>
return %0 : tensor<2x2x1xi32>
}
// CHECK-LITERAL: %cst = arith.constant dense<[[[3], [5]], [[9], [11]]]> : tensor<2x2x1xi32>
// CHECK-NOT: slice
// CHECK-LABEL: @slice_all_dims
func.func @slice_all_dims() -> tensor<1x2x1xi32> {
%cst_1 = arith.constant dense<[[[0, 1], [2, 3], [4, 5]], [[6, 7], [8, 9], [10, 11]]]> : tensor<2x3x2xi32>
%cst_2 = arith.constant dense<[1, 1, 1]> : tensor<3xi32>
%cst_3 = arith.constant dense<[1, 2, 1]> : tensor<3xi32>
%0 = "tfl.slice"(%cst_1, %cst_2, %cst_3) : (tensor<2x3x2xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x2x1xi32>
return %0 : tensor<1x2x1xi32>
}
// CHECK-LITERAL: %cst = arith.constant dense<[[[9], [11]]]> : tensor<1x2x1xi32>
// CHECK-NOT: slice
// CHECK-LABEL: @slice_some_dims_i64
func.func @slice_some_dims_i64() -> tensor<2x2x1xi32> {
%cst_1 = arith.constant dense<[[[0, 1], [2, 3], [4, 5]], [[6, 7], [8, 9], [10, 11]]]> : tensor<2x3x2xi32>
%cst_2 = arith.constant dense<[0, 1, 1]> : tensor<3xi64>
%cst_3 = arith.constant dense<[2, 2, 1]> : tensor<3xi64>
%0 = "tfl.slice"(%cst_1, %cst_2, %cst_3) : (tensor<2x3x2xi32>, tensor<3xi64>, tensor<3xi64>) -> tensor<2x2x1xi32>
return %0 : tensor<2x2x1xi32>
}
// CHECK-LITERAL: %cst = arith.constant dense<[[[3], [5]], [[9], [11]]]> : tensor<2x2x1xi32>
// CHECK-NOT: slice
// CHECK-LABEL: @slice_big_float
func.func @slice_big_float() -> tensor<1x1x1792x256xf32> {
%cst_1 = arith.constant dense<9.000000e+00> : tensor<2x1x1792x256xf32>
%cst_2 = arith.constant dense<[1, 0, 0, 0]> : tensor<4xi32>
%cst_3 = arith.constant dense<[1, 1, 1792, 256]> : tensor<4xi32>
%0 = "tfl.slice"(%cst_1, %cst_2, %cst_3) : (tensor<2x1x1792x256xf32>, tensor<4xi32>, tensor<4xi32>) -> tensor<1x1x1792x256xf32>
return %0 : tensor<1x1x1792x256xf32>
}
// CHECK-LITERAL: %cst = arith.constant dense<9.000000e+00> : tensor<1x1x1792x256xf32>
// CHECK-NOT: slice