| func.func @iota_dim0() { | |
| %result = "stablehlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<2x3xf32> | |
| check.expect_almost_eq_const(%result, dense<[ | |
| [0.0, 0.0, 0.0], | |
| [1.0, 1.0, 1.0]]> : tensor<2x3xf32>) : tensor<2x3xf32> | |
| return | |
| } | |
| func.func @iota_dim1() { | |
| %result = "stablehlo.iota"() {iota_dimension = 1 : i64} : () -> tensor<2x3xf32> | |
| check.expect_almost_eq_const(%result, dense<[ | |
| [0.0, 1.0, 2.0], | |
| [0.0, 1.0, 2.0]]> : tensor<2x3xf32>) : tensor<2x3xf32> | |
| return | |
| } |