blob: 5504b948aedf3004f7a897f6db993ad4a89f0d1e [file] [log] [blame]
// Note that they are stateless random generators, so they have fixed results.
func.func @rng_uniform_1d() {
%min = util.unfoldable_constant dense<-10.0> : tensor<f32>
%max = util.unfoldable_constant dense<10.0> : tensor<f32>
%shape = util.unfoldable_constant dense<[10]> : tensor<1xi32>
%res = "mhlo.rng_uniform"(%min, %max, %shape) : (tensor<f32>, tensor<f32>, tensor<1xi32>) -> tensor<10xf32>
check.expect_almost_eq_const(%res, dense<[
-9.99994, -4.8613, 0.277344, 5.41599, -9.44537, -4.30673, 0.831918, 5.97056, -8.8908, -3.75215
]> : tensor<10xf32>) : tensor<10xf32>
return
}
func.func @rng_uniform_2d() {
%min = util.unfoldable_constant dense<-10.0> : tensor<f32>
%max = util.unfoldable_constant dense<10.0> : tensor<f32>
%shape = util.unfoldable_constant dense<[3, 3]> : tensor<2xi32>
%res = "mhlo.rng_uniform"(%min, %max, %shape) : (tensor<f32>, tensor<f32>, tensor<2xi32>) -> tensor<3x3xf32>
check.expect_almost_eq_const(%res, dense<[
[6.55154, -8.30982, -3.17117],
[1.75741, 6.89606, -7.9653],
[-3.03671, 2.10193, 7.24057]]> : tensor<3x3xf32>) : tensor<3x3xf32>
return
}
func.func @rng_uniform_3d() {
%min = util.unfoldable_constant dense<-10.0> : tensor<f32>
%max = util.unfoldable_constant dense<10.0> : tensor<f32>
%shape = util.unfoldable_constant dense<[2, 2, 2]> : tensor<3xi32>
%res = "mhlo.rng_uniform"(%min, %max, %shape) : (tensor<f32>, tensor<f32>, tensor<3xi32>) -> tensor<2x2x2xf32>
check.expect_almost_eq_const(%res, dense<[
[[3.04814, 8.18679], [-1.74598, 3.39266]],
[[-6.91349, -1.77484], [8.29239, -6.56897]]]> : tensor<2x2x2xf32>) : tensor<2x2x2xf32>
return
}