return %1 : tensor<1x3xf32>
    }

    func.func private @composite_dot_general_fn(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> {
      %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
      return %0 : tensor<1x3xf32>
    }
  )mlir";

// CHECK-NEXT:      %22 = mhlo.dynamic_reshape %2, %21 : (tensor<2x?x3x4xf32>, tensor<4xi32>) -> tensor<2x?x3x4xf32>
// CHECK-NEXT:      %23 = "tfl.batch_matmul"(%12, %22) <{adj_x = false, adj_y = false, asymmetric_quantize_inputs = false}> : (tensor<2x?x2x3xf32>, tensor<2x?x3x4xf32>) -> tensor<2x?x2x4xf32>
// CHECK-NEXT:      %24 = "tfl.shape"(%arg0) : (tensor<2x?x2x3xf32>) -> tensor<4xi32>

    %8 = stablehlo.convert %7 : (tensor<1x4x2xi8>) -> tensor<1x4x2xf32>
    %9 = stablehlo.convert %2 : (tensor<2x3xi8>) -> tensor<2x3xf32>
    %10 = stablehlo.dot_general %8, %9, contracting_dims = [2] x [0] : (tensor<1x4x2xf32>, tensor<2x3xf32>) -> tensor<1x4x3xf32>
    %11 = stablehlo.convert %3 : (tensor<1x1x3xi32>) -> tensor<1x1x3xf32>

func.func @lower_einsum(%arg0: tensor<3x4x5xf32>, %arg1: tensor<3x5x6xf32>) -> tensor<3x4x6xf32> {
  %0 = "tf.Einsum"(%arg0, %arg1) {T = "tfdtype$DT_FLOAT", equation = "ijk,ikm->ijm"}: (tensor<3x4x5xf32>, tensor<3x5x6xf32>) -> tensor<3x4x6xf32>
  func.return %0 : tensor<3x4x6xf32>
}
// CHECK-LABEL: lower_einsum
// CHECK: "tf.BatchMatMulV2"(%arg0, %arg1) <{adj_x = false, adj_y = false}> : (tensor<3x4x5xf32>, tensor<3x5x6xf32>) -> tensor<3x4x6xf32>

// -----

}

// -----

func.func @trivial_dynamic_update_slice_wrong_update_shape(%arg0: tensor<2x7x14xf32>, %arg1: tensor<2x7x7xf32>) -> tensor<2x7x14xf32> {
  %0 = arith.constant dense<0> : tensor<3xi32>
  %1 = "tfl.dynamic_update_slice"(%arg0, %arg1, %0) : (tensor<2x7x14xf32>, tensor<2x7x7xf32>, tensor<3xi32>) -> tensor<2x7x14xf32>
  // CHECK: "tfl.dynamic_update_slice"
  func.return %1 : tensor<2x7x14xf32>
}