MegEngine/python_module/test/unit/module/test_tensor.py

# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
import copy

import numpy as np
import pytest

import megengine as mge
import megengine.functional as F
from megengine.core import Buffer, Graph, Parameter
from megengine.module import Conv2d
from megengine.test import assertTensorClose


def test_set_value():
    v0 = np.random.random((2, 3)).astype(np.float32)
    param = Parameter(v0)
    v1 = np.random.random((2, 3)).astype(np.float32)
    param.set_value(v1)
    assertTensorClose(param.numpy(), v1, max_err=5e-6)
    v2 = np.random.random((3, 3)).astype(np.float32)
    # TODO: add this
    # with pytest.raises(ValueError):
    #     param.set_value(v2)
    assertTensorClose(param.numpy(), v1, max_err=5e-6)


def test_fill():
    a = Buffer(np.zeros((2, 3), dtype=np.float32))
    a.fill(3)
    assertTensorClose(a.numpy(), np.full((2, 3), 3, dtype=np.float32))
    a.fill(124.568)
    assertTensorClose(a.numpy(), np.full((2, 3), 124.568, dtype=np.float32))


# TODO: remove or rewrite following test
# def test_attach():
#     p_ = np.random.random((2, 3)).astype(np.float32)

#     with Graph() as g:
#         g.set_option('eager_evaluation', False)
#         p = Parameter(p_)
#         v = p * 2
#         f = compile(v, None)

#     out, = f()
#     assertTensorClose(out, p_ * 2)

#     F.add_update(p, p)
#     out, = f()
#     assertTensorClose(out, p_ * 4)

# TODO: remove or rewrite following test
# def test_module_attach():
#     v = np.random.random((1, 3, 64, 64)).astype(np.float32)
#     net = Conv2d(3, 16, 3)

#     with Graph() as g:
#         g.set_option('eager_evaluation', False)

#         data0 = Input("data")
#         f = compile(net(data0), None)

#     out0, = f(data=v)

#     data1 = Input("data", value=v)
#     out1 = net(data1)

#     assertTensorClose(out0, out1.numpy())


def test_shape_warning():
    with Graph() as cg:
        cg.set_option("eager_evaluation", False)
        b = Buffer(np.ones((2, 3)).astype(np.float32))
        with pytest.warns(None) as record:
            print(b.shape)
        if len(record) != 0:
            raise ValueError(
                "Getting the shape of a constant Tensor should throw no Warning"
            )