%cst_2 = arith.constant dense<[4, 1024, 4096]> : tensor<3xi32>
  %0 = "tfl.fully_connected"(%arg0, %cst_0, %cst_1) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<4x1024x1024xbf16>, tensor<4096x1024xbf16>, none) -> tensor<4096x4096xbf16>
  %1 = "tfl.reshape"(%0, %cst_2) : (tensor<4096x4096xbf16>, tensor<3xi32>) -> tensor<4x1024x4096xbf16>
  func.return %1 : tensor<4x1024x4096xbf16>

  func.return %0 : tensor<2xi32>
}

// -----

// CHECK-LABEL: listFromTensor
func.func @listFromTensor(%tensor: tensor<3xi32>, %shape : tensor<?xi32>) -> tensor<!tf_type.variant<tensor<i32>>> {
  %0 = "tf.TensorListFromTensor"(%tensor, %shape) : (tensor<3xi32>, tensor<?xi32>) -> tensor<!tf_type.variant<tensor<i32>>>

// CHECK: return %1 : tensor<5x2x3x4xf32>

// -----

// CHECK-LABEL: pushTposeAfterAddSimpleWithFold
func.func @pushTposeAfterAddSimpleWithFold(%arg0: tensor<2x3xi32>) -> tensor<3x2xi32> {
  %perm = arith.constant dense<[1, 0]> : tensor<2xi32>
  %0 = "tfl.transpose"(%arg0, %perm) : (tensor<2x3xi32>, tensor<2xi32>) -> tensor<3x2xi32>

// CHECK:   %{{.*}} = "tfl.dequantize"(%{{.*}}) : (tensor<1x401408x!quant.uniform<u8:f32, 3.906250e-03>>) -> tensor<1x401408xf32>

  %cst = arith.constant dense<[1, 401408]> : tensor<2xi32>

// Note: Argument quantization sequence omitted.
// CHECK: stablehlo.dot_general %{{.+}}, %{{.+}}, contracting_dims = [1] x [0] : (tensor<1x1024xi8>, tensor<1024x3xi8>) -> tensor<1x3xi32>

// Note: Result dequantization sequence omitted.
// CHECK: return %{{.+}} : tensor<1x3xf32>

// -----

// Tests that a simple dot_general (lifted as a function) with CustomAggregators

  func.return %1 : tensor<1x1x257xi32>
  // CHECK: %[[CONST_DIM:.*]] = "tfl.pseudo_const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32>
  // CHECK: %[[RED:.*]] = "tfl.sum"(%arg0, %[[CONST_DIM]]) <{keep_dims = true}> : (tensor<1x16384x257xi32>, tensor<1xi32>) -> tensor<1x1x257xi32>
}

// CHECK-LABEL: BatchmatmulToReduceSumF32
// CHECK-NOT: "tfl.batch_matmul"