// CHECK-SAME:     scatter_dimension = 0
      //
      %1 = "tf.XlaReduceScatter"(%arg0, %cst_0, %cst) {reduce_op = "Add"} : (tensor<128x128xf32>, tensor<4x2xi32>, tensor<i32>) -> tensor<64x128xf32>
      func.return %1 : tensor<64x128xf32>
  }

  // CHECK-LABEL: func @tf_mod
  func.func @tf_mod(%arg1: tensor<2x2xf32>) -> tensor<2x2xf32> {
    %cst = "tf.Const"() {value = dense<7.000000e+00> : tensor<f32>} : () -> tensor<f32>

  func.return %4 : tensor<8x128xf32>

// CHECK-LABEL: SoftMaxWithNormalization
// CHECK: %[[RESULT:.*]] = "tfl.softmax"(%arg0) <{beta = 1.000000e+00 : f32}> : (tensor<8x128xf32>) -> tensor<8x128xf32>
// CHECK: return %[[RESULT]] : tensor<8x128xf32>
}

func.func @SoftMaxWithoutNormalization(%arg0: tensor<8x128xf32>) -> tensor<8x128xf32> {

  // CHECK-SAME:   slice_sizes = dense<[1, 4, 128]>
  // CHECK-SAME: (tensor<2x4x128xf32>, tensor<2x1xi32>) -> tensor<2x4x128xf32>
  %0 =  "tf.GatherNd"(%arg0, %arg1) {Tindices = i32, Tparams = i32, device = ""} : (tensor<2x4x128xf32>, tensor<2x1xi32>) -> tensor<2x4x128xf32>
  func.return %0 : tensor<2x4x128xf32>
}

//===----------------------------------------------------------------------===//

  func.return %1 : tensor<4x12xf32>
}

// CHECK-LABEL:   func @lowered_cumsum_trivial_attrs(
// CHECK-SAME:      %[[VAL_0:.*]]: tensor<4x12xf32>) -> tensor<4x12xf32> {
// CHECK-DAG:       %[[VAL_1:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<f32>}> : () -> tensor<f32>

// CHECK: return %[[v5]] : tensor<?x?x8xf32>
}

func.func @einsum_no_reshape(%arg0 : tensor<?x?x8x128xf32>, %arg1 : tensor<1x?x8x128xf32>) -> tensor<?x?x8x?xf32> {
  %0 = "tf.Einsum"(%arg0, %arg1) {device = "", equation = "BQNH,BTNH->BQNT"} : (tensor<?x?x8x128xf32>, tensor<1x?x8x128xf32>) -> tensor<?x?x8x?xf32>
  func.return %0 : tensor<?x?x8x?xf32>
// CHECK-LABEL: einsum_no_reshape

module {
    func.func @main(%arg0: tensor<8x128xf32>, %arg1: tensor<f32>) -> tensor<11x131xf32> {
      %0 = "tfl.custom"(%arg0, %arg1) {custom_code = "stablehlo.pad", custom_option = #tfl<const_bytes : "0x656467655F70616464696E675F6869676800020203656467655F70616464696E675F6C6F7700020100696E746572696F725F70616464696E6700020000033E2A17030103311E0B2C2C2C062401">} : (tensor<8x128xf32>, tensor<f32>) -> tensor<11x131xf32>

func.func @testConv2dShapeInferenceDynamic(%arg0: tensor<1x?x?x128xf32>, %arg1: tensor<128x3x3x128xf32>, %arg2: tensor<128xf32>) -> tensor<1x?x?x128xf32> {
  // CHECK: "tfl.conv_2d"(%arg0, %arg1, %arg2) <{dilation_h_factor = 2 : i32, dilation_w_factor = 2 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x?x?x128xf32>, tensor<128x3x3x128xf32>, tensor<128xf32>) -> tensor<1x?x?x128xf32>

  func.return %1 : tensor<16x128xf32>
  // CHECK: return %0 : tensor<16x128xf32>
}

// CHECK-LABEL: FuseTransposeFCLhsToBatchMatmul

  %w = arith.constant dense<127.0> : tensor<4x12xf32>
  %b = arith.constant dense<0.0> : tensor<4xf32>
  %fc = "tfl.fully_connected"(%arg0, %w, %b) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<1x224x224x3xf32>, tensor<4x12xf32>, tensor<4xf32>) -> tensor<1x112x112x4xf32>
  func.return %fc : tensor<1x112x112x4xf32>

// CHECK:           %[[VAL_15:.*]] = "tfl.reshape"(%[[VAL_14]], %[[VAL_5]]) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<384x128xf32>, tensor<3xi32>) -> tensor<1x384x128xf32>