}

// -----

// CHECK-LABEL: testLiftDotConcatLHSAndRHS
func.func @testLiftDotConcatLHSAndRHS(%arg0: tensor<1x72x128xf32>, %arg1: tensor<1x128x72xf32>, %arg2: tensor<1x72x128xf32>, %arg3: tensor<1x128x72xf32>, %arg4: tensor<1x72x128xf32>, %arg5: tensor<1x128x72xf32>, %arg6: tensor<1x72x128xf32>, %arg7: tensor<1x128x72xf32>) -> tensor<4x72x72xf32> {
  %0 = "mhlo.dot_general"(%arg0, %arg1) {
    dot_dimension_numbers = #mhlo.dot<

    // CHECK: mhlo.select
    %2 = "tf.MatrixDiagPartV3"(%arg0, %1, %0) {
        T = i32, align = "RIGHT_LEFT"
    } : (tensor<7x140x128xi32>, tensor<2xi32>, tensor<i32>) -> tensor<7x22x128xi32>
    func.return %2: tensor<7x22x128xi32>
  }

  // CHECK-LABEL: func @xla_call_module
  func.func @xla_call_module(%arg0: tensor<f32>) -> tensor<*xf32> {
    // Equivalent to the following:
    //
    // module @jit_sin {