tests/e2e/stablehlo_ops/broadcast.mlir

func.func @broadcast_2D_3D() {
  %input = util.unfoldable_constant dense<[[1, 2, 3, 4],
                                           [5, 6, 7, 8]]> : tensor<2x4xi32>
  %result = "stablehlo.broadcast"(%input) {broadcast_sizes = array<i64: 3>} : (tensor<2x4xi32>) -> tensor<3x2x4xi32>
  check.expect_eq_const(%result, dense<[
      [[1, 2, 3, 4], [5, 6, 7, 8]],
      [[1, 2, 3, 4], [5, 6, 7, 8]],
      [[1, 2, 3, 4], [5, 6, 7, 8]]]> : tensor<3x2x4xi32>) : tensor<3x2x4xi32>
  return
}

func.func @broadcast_3D_scalar() {
  %input = util.unfoldable_constant dense<42> : tensor<i32>
  %result = "stablehlo.broadcast"(%input) {broadcast_sizes = array<i64: 3, 2, 4>} : (tensor<i32>) -> tensor<3x2x4xi32>
  check.expect_eq_const(%result, dense<[
      [[42, 42, 42, 42], [42, 42, 42, 42]],
      [[42, 42, 42, 42], [42, 42, 42, 42]],
      [[42, 42, 42, 42], [42, 42, 42, 42]]]> : tensor<3x2x4xi32>) : tensor<3x2x4xi32>
  return
}