forked from huawei/mindspore2022
221 lines
9.3 KiB
Python
221 lines
9.3 KiB
Python
# Copyright 2022 Huawei Technologies Co., Ltd
|
|
#
|
|
# Licensed 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.
|
|
# ============================================================================
|
|
"""test vmap in graph mode"""
|
|
import pytest
|
|
import mindspore.nn as nn
|
|
import mindspore.context as context
|
|
import mindspore.ops.operations as P
|
|
from mindspore import Tensor
|
|
from mindspore import dtype as mstype
|
|
from mindspore.ops.functional import vmap
|
|
|
|
context.set_context(mode=context.GRAPH_MODE)
|
|
|
|
|
|
class ThreeInputsTwoOutputsNet(nn.Cell):
|
|
def construct(self, x, y, z):
|
|
return x + y, z
|
|
|
|
|
|
def test_lambda_fn():
|
|
"""
|
|
Feature: vmap
|
|
Description: The first argument of `vmap` is a lambda function.
|
|
Expectation: throw TypeError:"Parse Lambda Function Fail. Node type must be Lambda, but got Call."
|
|
"""
|
|
x_hat = Tensor([[1, 2, 3], [4, 5, 6]], mstype.float32)
|
|
y_hat = Tensor([[1, 2, 3], [4, 5, 6]], mstype.float32)
|
|
z_hat = 1
|
|
with pytest.raises(TypeError) as ex:
|
|
vmap(lambda x, y, z: x + y + z, in_axes=(1, 1, None), out_axes=0)(x_hat, y_hat, z_hat)
|
|
assert "Parse Lambda Function Fail. Node type must be Lambda, but got Call." in str(ex.value)
|
|
|
|
|
|
def test_single_op():
|
|
"""
|
|
Feature: vmap
|
|
Description: The first argument of `vmap` is a single primitive.
|
|
Expectation: throw RuntimeError:"'VmapOperation' arg0 Prim: S-Prim-Add cast to 'FuncGraphAbstractClosure' failed."
|
|
"""
|
|
x_hat = Tensor([[1, 2, 3], [4, 5, 6]], mstype.float32)
|
|
y_hat = Tensor([[1, 2, 3], [4, 5, 6]], mstype.float32)
|
|
with pytest.raises(RuntimeError) as ex:
|
|
vmap(P.Add(), in_axes=(1, 1), out_axes=0)(x_hat, y_hat)
|
|
assert "'VmapOperation' arg0 Prim: S-Prim-Add cast to 'FuncGraphAbstractClosure' failed." in str(ex.value)
|
|
|
|
|
|
def test_none_in_axes():
|
|
"""
|
|
Feature: vmap
|
|
Description: The `in_axis` argument of `vmap` is a single None, and it's invalid when apply `vmap`.
|
|
Expectation: throw RuntimeError:"The 'in_axes' of 'vmap' cannot be a single None."
|
|
"""
|
|
x_hat = Tensor([[1, 2, 3], [4, 5, 6]], mstype.float32)
|
|
y_hat = Tensor([[1, 2, 3], [4, 5, 6]], mstype.float32)
|
|
z_hat = 1
|
|
with pytest.raises(RuntimeError) as ex:
|
|
vmap(ThreeInputsTwoOutputsNet(), in_axes=None, out_axes=0)(x_hat, y_hat, z_hat)
|
|
assert "The 'in_axes' of 'vmap' cannot be a single None." in str(ex.value)
|
|
|
|
|
|
def test_none_out_axes():
|
|
"""
|
|
Feature: vmap
|
|
Description: The `out_axes` argument of `vmap` is a nested None, and it's invalid when apply `vmap`.
|
|
Expectation: throw RuntimeError:"The 'out_axes' of 'vmap' cannot be all None, but got
|
|
(None, None, None, (None, None))."
|
|
"""
|
|
x_hat = Tensor([[1, 2, 3], [4, 5, 6]], mstype.float32)
|
|
y_hat = Tensor([[1, 2, 3], [4, 5, 6]], mstype.float32)
|
|
z_hat = 1
|
|
with pytest.raises(RuntimeError) as ex:
|
|
vmap(ThreeInputsTwoOutputsNet(), in_axes=(1, 1, None),
|
|
out_axes=(None, None, None, (None, None)))(x_hat, y_hat, z_hat)
|
|
assert "The 'out_axes' of 'vmap' cannot be all None, but got (None, None, None, (None, None))." in str(ex.value)
|
|
|
|
|
|
def test_mismatch_out_axes():
|
|
"""
|
|
Feature: vmap
|
|
Description: The `out_axes` of `vmap` sets to (0, 0, 0), but the outputs of `fn` is x + y, z.
|
|
Expectation: throw RuntimeError:"The size of vmap's 'out_axes' should be equal to the number of results of 'fn': 2,
|
|
but got size: 3."
|
|
"""
|
|
x_hat = Tensor([[1, 2, 3], [4, 5, 6]], mstype.float32)
|
|
y_hat = Tensor([[1, 2, 3], [4, 5, 6]], mstype.float32)
|
|
z_hat = 1
|
|
with pytest.raises(RuntimeError) as ex:
|
|
vmap(ThreeInputsTwoOutputsNet(), in_axes=(1, 1, None), out_axes=(0, 0, 0))(x_hat, y_hat, z_hat)
|
|
assert "The size of vmap's 'out_axes' should be equal to the number of results of 'fn': 2, but got size: 3." \
|
|
in str(ex.value)
|
|
|
|
|
|
def test_axis_type():
|
|
"""
|
|
Feature: vmap
|
|
Description: The `in_axes` of `vmap` contains elements of Float type.
|
|
Expectation: throw RuntimeError:"The axis in vmap's 'in_axes' should be a None or a scalar of type Int64Imm,
|
|
but got a 1."
|
|
"""
|
|
x_hat = Tensor([[1, 2, 3], [4, 5, 6]], mstype.float32)
|
|
y_hat = Tensor([[1, 2, 3], [4, 5, 6]], mstype.float32)
|
|
z_hat = 1
|
|
with pytest.raises(RuntimeError) as ex:
|
|
vmap(ThreeInputsTwoOutputsNet(), in_axes=(1., 1., None), out_axes=0)(x_hat, y_hat, z_hat)
|
|
assert "The axis in vmap's 'in_axes' should be a None or a scalar of type Int64Imm, but got a 1." in str(ex.value)
|
|
|
|
|
|
def test_axis_out_of_bounds():
|
|
"""
|
|
Feature: vmap
|
|
Description: The dimension of X is 2, but the corresponding axis -3 is set.
|
|
Expectation: throw RuntimeError:"The axis: -3 in 'in_axes' is out of bounds for array of dimension [-2,2)."
|
|
"""
|
|
x_hat = Tensor([[1, 2, 3], [4, 5, 6]], mstype.float32)
|
|
y_hat = Tensor([[1, 2, 3], [4, 5, 6]], mstype.float32)
|
|
z_hat = 1
|
|
with pytest.raises(RuntimeError) as ex:
|
|
vmap(ThreeInputsTwoOutputsNet(), in_axes=(-3, 2, None), out_axes=0)(x_hat, y_hat, z_hat)
|
|
assert "The axis: -3 in 'in_axes' is out of bounds for array of dimension [-2,2)." in str(ex.value)
|
|
|
|
|
|
def test_mismatch_none_axis():
|
|
"""
|
|
Feature: vmap
|
|
Description: The source axis of the first output of `fn` is non-None, but the `out_axes` for that is None,
|
|
it's invalid when apply `vmap`.
|
|
Expectation: throw RuntimeError:"It is invalid that source is not None and dst is None."
|
|
"""
|
|
x_hat = Tensor([[1, 2, 3], [4, 5, 6]], mstype.float32)
|
|
y_hat = Tensor([[1, 2, 3], [4, 5, 6]], mstype.float32)
|
|
z_hat = 1
|
|
with pytest.raises(RuntimeError) as ex:
|
|
vmap(ThreeInputsTwoOutputsNet(), in_axes=(1, 1, None), out_axes=(None, 0))(x_hat, y_hat, z_hat)
|
|
assert "It is invalid that source is not None and dst is None." in str(ex.value)
|
|
|
|
|
|
def test_mismatch_parameters_number():
|
|
"""
|
|
Feature: vmap
|
|
Description: The arguments of the cell is (x, y, z), but the arguments of vmap-ed function is (x_hat, y_hat).
|
|
Expectation: throw TypeError:"The parameters number of the function is 3, but the number of provided arguments
|
|
is 2."
|
|
"""
|
|
x_hat = Tensor([[1, 2, 3], [4, 5, 6]], mstype.float32)
|
|
y_hat = Tensor([[1, 2, 3], [4, 5, 6]], mstype.float32)
|
|
with pytest.raises(TypeError) as ex:
|
|
vmap(ThreeInputsTwoOutputsNet(), in_axes=(1, 1, None), out_axes=0)(x_hat, y_hat)
|
|
assert "The parameters number of the function is 3, but the number of provided arguments is 2." in str(ex.value)
|
|
|
|
|
|
def test_mismatch_axis_size():
|
|
"""
|
|
Feature: vmap
|
|
Description: The `axis_size` of X is 3, and the `axis_size` of Y is 2, they are not equal, vmap needs to ensure
|
|
that the `axis_size` of all parameters are uniform.
|
|
Expectation: throw RuntimeError:"The 'axis_size' of each argument in the scope of 'vmap' should be equal,
|
|
but got 3 and 2."
|
|
"""
|
|
x_hat = Tensor([[1, 2, 3], [4, 5, 6]], mstype.float32)
|
|
y_hat = Tensor([[1, 2, 3], [4, 5, 6]], mstype.float32)
|
|
z_hat = 1
|
|
with pytest.raises(RuntimeError) as ex:
|
|
vmap(ThreeInputsTwoOutputsNet(), in_axes=(1, 0, None), out_axes=0)(x_hat, y_hat, z_hat)
|
|
assert "The 'axis_size' of each argument in the scope of 'vmap' should be equal, but got 3 and 2." in str(ex.value)
|
|
|
|
|
|
def test_vmap_non_input():
|
|
"""
|
|
Feature: vmap
|
|
Description: The arguments of the cell is empty, it's invalid when apply `vmap`.
|
|
Expectation: throw RuntimeError:"Failed to get 'axis_size' within the scope of vmap."
|
|
"""
|
|
class NonInputSingleOutputNet(nn.Cell):
|
|
def construct(self):
|
|
return 1
|
|
|
|
with pytest.raises(RuntimeError) as ex:
|
|
vmap(NonInputSingleOutputNet())()
|
|
assert "Failed to get 'axis_size' within the scope of vmap." in str(ex.value)
|
|
|
|
|
|
def test_non_fn():
|
|
"""
|
|
Feature: vmap
|
|
Description: The first argument of `vmap` not provided, which is required positional argument.
|
|
Expectation: throw TypeError:"vmap() missing 1 required positional argument: 'fn'"
|
|
"""
|
|
x_hat = Tensor([[1, 2, 3], [4, 5, 6]], mstype.float32)
|
|
y_hat = Tensor([[1, 2, 3], [4, 5, 6]], mstype.float32)
|
|
z_hat = 1
|
|
with pytest.raises(TypeError) as ex:
|
|
vmap(in_axes=(1, 1, None), out_axes=0)(x_hat, y_hat, z_hat)
|
|
assert "vmap() missing 1 required positional argument: 'fn'" in str(ex.value)
|
|
|
|
|
|
def test_scalar_with_non_zero_axis():
|
|
"""
|
|
Feature: vmap
|
|
Description: The second output of `fn` is a scalar with source axis None, but get a destination axis 1, and it's
|
|
invalid when apply `vmap`.
|
|
Expectation: throw RuntimeError:"The axis: 1 in 'out_axes' is out of bounds for array of dimension [-1,1)."
|
|
"""
|
|
x_hat = Tensor([[1, 2, 3], [4, 5, 6]], mstype.float32)
|
|
y_hat = Tensor([[1, 2, 3], [4, 5, 6]], mstype.float32)
|
|
z_hat = 1
|
|
with pytest.raises(RuntimeError) as ex:
|
|
vmap(ThreeInputsTwoOutputsNet(), in_axes=(1, 1, None), out_axes=(0, 1))(x_hat, y_hat, z_hat)
|
|
assert "The axis: 1 in 'out_axes' is out of bounds for array of dimension [-1,1)." in str(ex.value)
|