func.func @iota_dim0() { | |
%result = "mhlo.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 = "mhlo.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 | |
} |