// 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 {