func.func @scan_1d_dim0_inclusive_sum() {
  %input = util.unfoldable_constant dense<[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]> : tensor<6xf32>

  %init = linalg.init_tensor [6] : tensor<6xf32>
  %t0 = util.unfoldable_constant dense<0.0> : tensor<f32>
  %0:2 = iree_linalg_ext.scan
         dimension(0) inclusive(true)
         ins(%input : tensor<6xf32>)
         outs(%init, %t0 : tensor<6xf32>, tensor<f32>) {
           ^bb0(%arg0 : f32, %arg1 : f32):
             %sum = arith.addf %arg0, %arg1 : f32
             iree_linalg_ext.yield %sum : f32
         } -> tensor<6xf32>, tensor<f32>

  check.expect_almost_eq_const(
      %0#0,
      dense<[1.0, 3.0, 6.0, 10.0, 15.0, 21.0]> : tensor<6xf32>
  ) : tensor<6xf32>

  check.expect_almost_eq_const(
      %0#1,
      dense<21.0> : tensor<f32>
  ) : tensor<f32>

  return
}

func.func @scan_1d_dim0_exclusive_sum() {
  %input = util.unfoldable_constant dense<[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]> : tensor<6xf32>

  %init = linalg.init_tensor [6] : tensor<6xf32>
  %t0 = util.unfoldable_constant dense<10.0> : tensor<f32>
  %0:2 = iree_linalg_ext.scan
         dimension(0) inclusive(false)
         ins(%input : tensor<6xf32>)
         outs(%init, %t0 : tensor<6xf32>, tensor<f32>) {
           ^bb0(%arg0 : f32, %arg1 : f32):
             %sum = arith.addf %arg0, %arg1 : f32
             iree_linalg_ext.yield %sum : f32
         } -> tensor<6xf32>, tensor<f32>

  check.expect_almost_eq_const(
      %0#0,
      dense<[10.0, 11.0, 13.0, 16.0, 20.0, 25.0]> : tensor<6xf32>
  ) : tensor<6xf32>

  check.expect_almost_eq_const(
      %0#1,
      dense<25.0> : tensor<f32>
  ) : tensor<f32>

  return
}

func.func @scan_1d_dim0_inclusive_mul() {
  %input = util.unfoldable_constant dense<[1, 2, 3, 4, 5, 6]> : tensor<6xi32>

  %init = linalg.init_tensor [6] : tensor<6xi32>
  %t0 = util.unfoldable_constant dense<1> : tensor<i32>
  %0:2 = iree_linalg_ext.scan
         dimension(0) inclusive(true)
         ins(%input : tensor<6xi32>)
         outs(%init, %t0 : tensor<6xi32>, tensor<i32>) {
           ^bb0(%arg0 : i32, %arg1 : i32):
             %sum = arith.muli %arg0, %arg1 : i32
             iree_linalg_ext.yield %sum : i32
         } -> tensor<6xi32>, tensor<i32>

  check.expect_eq_const(
      %0#0,
      dense<[1, 2, 6, 24, 120, 720]> : tensor<6xi32>
  ) : tensor<6xi32>

  check.expect_eq_const(
      %0#1,
      dense<720> : tensor<i32>
  ) : tensor<i32>

  return
}

func.func @scan_2d_dim0_inclusive_sum() {
  %input = util.unfoldable_constant dense<[[1, 2, 3],
                                           [4, 5, 6]]> : tensor<2x3xi32>

  %init = linalg.init_tensor [2, 3] : tensor<2x3xi32>
  %t0 = util.unfoldable_constant dense<[0, 0, 0]> : tensor<3xi32>
  %0:2 = iree_linalg_ext.scan
         dimension(0) inclusive(true)
         ins(%input : tensor<2x3xi32>)
         outs(%init, %t0 : tensor<2x3xi32>, tensor<3xi32>) {
           ^bb0(%arg0 : i32, %arg1 : i32):
             %sum = arith.addi %arg0, %arg1 : i32
             iree_linalg_ext.yield %sum : i32
         } -> tensor<2x3xi32>, tensor<3xi32>

  check.expect_eq_const(
      %0#0,
      dense<[[1, 2, 3], [5, 7, 9]]> : tensor<2x3xi32>
  ) : tensor<2x3xi32>

  check.expect_eq_const(
      %0#1,
      dense<[5, 7, 9]> : tensor<3xi32>
  ) : tensor<3xi32>

  return
}

func.func @scan_2d_dim1_inclusive_sum() {
  %input = util.unfoldable_constant dense<[[1, 2, 3],
                                           [4, 5, 6]]> : tensor<2x3xi32>

  %init = linalg.init_tensor [2, 3] : tensor<2x3xi32>
  %t0 = util.unfoldable_constant dense<[0, 0]> : tensor<2xi32>
  %0:2 = iree_linalg_ext.scan
         dimension(1) inclusive(true)
         ins(%input : tensor<2x3xi32>)
         outs(%init, %t0 : tensor<2x3xi32>, tensor<2xi32>) {
           ^bb0(%arg0 : i32, %arg1 : i32):
             %sum = arith.addi %arg0, %arg1 : i32
             iree_linalg_ext.yield %sum : i32
         } -> tensor<2x3xi32>, tensor<2xi32>

  check.expect_eq_const(
      %0#0,
      dense<[[1, 3, 6], [4, 9, 15]]> : tensor<2x3xi32>
  ) : tensor<2x3xi32>

  check.expect_eq_const(
      %0#1,
      dense<[6, 15]> : tensor<2xi32>
  ) : tensor<2xi32>

  return
}