test_random.py 3.07 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
17
import tvm
18
from tvm import te
19 20 21 22 23 24 25
import numpy as np
from tvm.contrib import random

def test_randint():
    m = 1024
    n = 1024
    A = random.randint(-127, 128, size=(m, n), dtype='int32')
26
    s = te.create_schedule(A.op)
27 28

    def verify(target="llvm"):
29
        if not tvm.runtime.enabled(target):
30 31 32
            print("skip because %s is not enabled..." % target)
            return
        if not tvm.get_global_func("tvm.contrib.random.randint", True):
33
            print("skip because extern function is not available")
34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49
            return
        ctx = tvm.cpu(0)
        f = tvm.build(s, [A], target)
        a = tvm.nd.array(np.zeros((m, n), dtype=A.dtype), ctx)
        f(a)
        na = a.asnumpy()
        assert abs(np.mean(na)) < 0.2
        assert np.min(na) == -127
        assert np.max(na) == 127
    verify()


def test_uniform():
    m = 1024
    n = 1024
    A = random.uniform(0, 1, size=(m, n))
50
    s = te.create_schedule(A.op)
51 52

    def verify(target="llvm"):
53
        if not tvm.runtime.enabled(target):
54 55 56
            print("skip because %s is not enabled..." % target)
            return
        if not tvm.get_global_func("tvm.contrib.random.uniform", True):
57
            print("skip because extern function is not available")
58 59 60 61 62 63 64 65 66 67 68 69
            return
        ctx = tvm.cpu(0)
        f = tvm.build(s, [A], target)
        a = tvm.nd.array(np.zeros((m, n), dtype=A.dtype), ctx)
        f(a)
        na = a.asnumpy()
        assert abs(np.mean(na) - 0.5) < 1e-2
        assert abs(np.min(na) - 0.0) < 1e-3
        assert abs(np.max(na) - 1.0) < 1e-3
    verify()


70 71 72 73
def test_normal():
    m = 1024
    n = 1024
    A = random.normal(3, 4, size=(m, n))
74
    s = te.create_schedule(A.op)
75 76

    def verify(target="llvm"):
77
        if not tvm.runtime.enabled(target):
78 79 80
            print("skip because %s is not enabled..." % target)
            return
        if not tvm.get_global_func("tvm.contrib.random.normal", True):
81
            print("skip because extern function is not available")
82 83 84 85 86 87 88 89 90 91 92
            return
        ctx = tvm.cpu(0)
        f = tvm.build(s, [A], target)
        a = tvm.nd.array(np.zeros((m, n), dtype=A.dtype), ctx)
        f(a)
        na = a.asnumpy()
        assert abs(np.mean(na) - 3) < 1e-2
        assert abs(np.std(na) - 4) < 1e-2
    verify()


93 94 95
if __name__ == "__main__":
    test_randint()
    test_uniform()
96
    test_normal()