@@ -9,6 +9,8 @@ | |||||
import os | import os | ||||
from ._imperative_rt.core2 import set_cpp_use_symbolic_shape | |||||
_use_symbolic_shape = False | _use_symbolic_shape = False | ||||
if os.environ.get("MEGENGINE_USE_SYMBOLIC_SHAPE"): | if os.environ.get("MEGENGINE_USE_SYMBOLIC_SHAPE"): | ||||
_use_symbolic_shape = True | _use_symbolic_shape = True | ||||
@@ -25,3 +27,6 @@ def set_symbolic_shape(option: bool): | |||||
_org = _use_symbolic_shape | _org = _use_symbolic_shape | ||||
_use_symbolic_shape = option | _use_symbolic_shape = option | ||||
return _org | return _org | ||||
set_cpp_use_symbolic_shape(use_symbolic_shape) |
@@ -22,12 +22,12 @@ from .._imperative_rt.core2 import ( | |||||
astype_cpp, | astype_cpp, | ||||
broadcast_cpp, | broadcast_cpp, | ||||
dtype_promotion, | dtype_promotion, | ||||
getitem_cpp, | |||||
) | ) | ||||
from .._imperative_rt.core2 import reduce_to_scalar as _reduce_to_scalar | from .._imperative_rt.core2 import reduce_to_scalar as _reduce_to_scalar | ||||
from .._imperative_rt.core2 import reshape_cpp, squeeze_cpp, transpose_cpp | |||||
from .._imperative_rt.core2 import reshape_cpp, setitem_cpp, squeeze_cpp, transpose_cpp | |||||
from ..ops import builtin | from ..ops import builtin | ||||
from . import amp | from . import amp | ||||
from .indexing import getitem, setitem | |||||
from .utils import _normalize_axis, astensor1d, cast_tensors, make_shape_tuple, subgraph | from .utils import _normalize_axis, astensor1d, cast_tensors, make_shape_tuple, subgraph | ||||
_ElwMod = builtin.Elemwise.Mode | _ElwMod = builtin.Elemwise.Mode | ||||
@@ -544,11 +544,11 @@ class ArrayMethodMixin(abc.ABC): | |||||
yield self[i] | yield self[i] | ||||
def __getitem__(self, index): | def __getitem__(self, index): | ||||
return getitem(self, index) | |||||
return getitem_cpp(self, index) | |||||
def __setitem__(self, index, value): | def __setitem__(self, index, value): | ||||
if index is not Ellipsis: | if index is not Ellipsis: | ||||
value = setitem(self, index, value) | |||||
value = setitem_cpp(self, index, value) | |||||
self._reset(value) | self._reset(value) | ||||
__contains__ = _todo | __contains__ = _todo | ||||
@@ -1,28 +0,0 @@ | |||||
# -*- coding: utf-8 -*- | |||||
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") | |||||
# | |||||
# Copyright (c) 2014-2021 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. | |||||
from .._imperative_rt.core2 import ( | |||||
getitem_cpp, | |||||
set_cpp_astensor1d, | |||||
set_cpp_use_symbolic_shape, | |||||
setitem_cpp, | |||||
) | |||||
from .._trace_option import use_symbolic_shape | |||||
from .utils import astensor1d | |||||
def getitem(tensor, index): | |||||
return getitem_cpp(tensor, index) | |||||
def setitem(tensor, index, value): | |||||
return setitem_cpp(tensor, index, value) | |||||
set_cpp_use_symbolic_shape(use_symbolic_shape) | |||||
set_cpp_astensor1d(astensor1d) |
@@ -20,6 +20,7 @@ from .._imperative_rt.core2 import ( | |||||
_get_convert_inputs, | _get_convert_inputs, | ||||
_set_convert_inputs, | _set_convert_inputs, | ||||
apply, | apply, | ||||
astensor1d_cpp, | |||||
astype_cpp, | astype_cpp, | ||||
convert_inputs_cpp, | convert_inputs_cpp, | ||||
convert_single_value_cpp, | convert_single_value_cpp, | ||||
@@ -50,14 +51,6 @@ def set_convert_inputs(flag): | |||||
return _set_convert_inputs(flag) | return _set_convert_inputs(flag) | ||||
def concatenate(inputs, axis=0, *, device=None): | |||||
inputs = convert_inputs(*inputs) | |||||
if device is None: | |||||
device = get_device(inputs) | |||||
(result,) = apply(builtin.Concat(axis=axis, comp_node=device), *inputs) | |||||
return result | |||||
def convert_single_value(v, *, dtype=None, device=None): | def convert_single_value(v, *, dtype=None, device=None): | ||||
return convert_single_value_cpp(v, dtype, device) | return convert_single_value_cpp(v, dtype, device) | ||||
@@ -104,34 +97,7 @@ def astensor1d(x, *reference, dtype=None, device=None): | |||||
* numpy array | * numpy array | ||||
* tensor (returned as is, regardless of dtype and device) | * tensor (returned as is, regardless of dtype and device) | ||||
""" | """ | ||||
try: | |||||
ndim = x.ndim | |||||
except AttributeError: | |||||
pass | |||||
except ValueError: | |||||
if dtype is not None and dtype != x.dtype: | |||||
x = astype_cpp(x, dtype) | |||||
if device is not None: | |||||
cn = as_device(device).to_c() | |||||
(x,) = apply(builtin.Copy(comp_node=cn), x) | |||||
return x | |||||
else: | |||||
if ndim != 0 and ndim != 1: | |||||
raise ValueError("ndim != 1 or 0, get : %d" % ndim) | |||||
if not isinstance(x, (Tensor, SymbolVar)): | |||||
x = Const(x, dtype, device, reference) | |||||
return x | |||||
if not isinstance(x, collections.abc.Sequence): | |||||
raise TypeError | |||||
if any(isinstance(i, (Tensor, SymbolVar)) for i in x): | |||||
x = concatenate(x, device=device) if len(x) > 1 else x[0] | |||||
if dtype is not None: | |||||
x = astype_cpp(x, dtype) | |||||
return x | |||||
x = Const(x, dtype, device, reference) | |||||
return x | |||||
return astensor1d_cpp(x, dtype, device, reference) | |||||
def _normalize_axis( | def _normalize_axis( | ||||
@@ -104,13 +104,12 @@ struct SymbolVarContext { | |||||
interpreter::Interpreter::Channel* interpreter_for_py = nullptr; | interpreter::Interpreter::Channel* interpreter_for_py = nullptr; | ||||
PyTypeObject* py_tensor_type = nullptr; | PyTypeObject* py_tensor_type = nullptr; | ||||
PyObject *cpp_use_symbolic_shape, *cpp_astensor1d; | |||||
PyObject* cpp_use_symbolic_shape; | |||||
#define REGISTE_APPLY_FUNC(mode) \ | #define REGISTE_APPLY_FUNC(mode) \ | ||||
void set_##mode(py::object pyf) { mode = pyf.ptr(); } | void set_##mode(py::object pyf) { mode = pyf.ptr(); } | ||||
REGISTE_APPLY_FUNC(cpp_use_symbolic_shape) | REGISTE_APPLY_FUNC(cpp_use_symbolic_shape) | ||||
REGISTE_APPLY_FUNC(cpp_astensor1d) | |||||
#undef REGISTE_APPLY_FUNC | #undef REGISTE_APPLY_FUNC | ||||
@@ -426,6 +425,7 @@ WRAP_FUNC_PY35(Const); | |||||
WRAP_FUNC_PY35(astype_cpp); | WRAP_FUNC_PY35(astype_cpp); | ||||
WRAP_FUNC_PY35(convert_single_value_cpp); | WRAP_FUNC_PY35(convert_single_value_cpp); | ||||
WRAP_FUNC_PY35(convert_inputs_cpp); | WRAP_FUNC_PY35(convert_inputs_cpp); | ||||
WRAP_FUNC_PY35(astensor1d_cpp); | |||||
#undef WRAP_FUNC_PY35 | #undef WRAP_FUNC_PY35 | ||||
#define MGE_PY_INTERFACE(NAME, FUNC) \ | #define MGE_PY_INTERFACE(NAME, FUNC) \ | ||||
{ #NAME, (PyCFunction)py35_##FUNC, METH_VARARGS, nullptr } | { #NAME, (PyCFunction)py35_##FUNC, METH_VARARGS, nullptr } | ||||
@@ -568,6 +568,7 @@ void init_tensor(py::module m) { | |||||
MGE_PY_INTERFACE(astype_cpp, astype_cpp), | MGE_PY_INTERFACE(astype_cpp, astype_cpp), | ||||
MGE_PY_INTERFACE(convert_single_value_cpp, convert_single_value_cpp), | MGE_PY_INTERFACE(convert_single_value_cpp, convert_single_value_cpp), | ||||
MGE_PY_INTERFACE(convert_inputs_cpp, convert_inputs_cpp), | MGE_PY_INTERFACE(convert_inputs_cpp, convert_inputs_cpp), | ||||
MGE_PY_INTERFACE(astensor1d_cpp, astensor1d_cpp), | |||||
{nullptr, nullptr, 0, nullptr}}; | {nullptr, nullptr, 0, nullptr}}; | ||||
for (auto&& def : method_defs) { | for (auto&& def : method_defs) { | ||||
if (def.ml_meth != nullptr) { | if (def.ml_meth != nullptr) { | ||||
@@ -957,8 +958,6 @@ void init_tensor(py::module m) { | |||||
m.def("set_cpp_use_symbolic_shape", &set_cpp_use_symbolic_shape); | m.def("set_cpp_use_symbolic_shape", &set_cpp_use_symbolic_shape); | ||||
m.def("set_cpp_astensor1d", &set_cpp_astensor1d); | |||||
m.def("set_module_tracing", [=] { get_module_trace()->enable(); }); | m.def("set_module_tracing", [=] { get_module_trace()->enable(); }); | ||||
m.def("unset_module_tracing", [=] { get_module_trace()->disable(); }); | m.def("unset_module_tracing", [=] { get_module_trace()->disable(); }); | ||||
@@ -32,4 +32,6 @@ PyObject* convert_single_value_cpp(PyObject* self, PyObject* const* args, size_t | |||||
PyObject* convert_inputs_cpp(PyObject* self, PyObject* const* args, size_t nargs); | PyObject* convert_inputs_cpp(PyObject* self, PyObject* const* args, size_t nargs); | ||||
PyObject* astensor1d_cpp(PyObject* self, PyObject* const* args, size_t nargs); | |||||
} // namespace mgb::imperative::python | } // namespace mgb::imperative::python |
@@ -511,6 +511,20 @@ def test_advance_indexing_with_bool(test_varnode): | |||||
network = Network() | network = Network() | ||||
else: | else: | ||||
network = None | network = None | ||||
a = np.array([[True, False], [False, True]]) | |||||
b = np.array([1]) | |||||
aa = make_tensor(a, network) | |||||
bb = make_tensor(b, network) | |||||
np.testing.assert_equal(a[b], get_value(aa[bb])) | |||||
b = np.array([[True, True], [False, True]]) | |||||
bb = make_tensor(b, network) | |||||
np.testing.assert_equal(a[b], get_value(aa[bb])) | |||||
if not test_varnode: | |||||
a[b] = False | |||||
aa[bb] = False | |||||
np.testing.assert_equal(a, get_value(aa)) | |||||
a = np.arange(9).reshape(3, 3).astype(np.float32) | a = np.arange(9).reshape(3, 3).astype(np.float32) | ||||
b = np.array([1, 2, 3]) | b = np.array([1, 2, 3]) | ||||
c = np.array([1, 2, 3]) | c = np.array([1, 2, 3]) | ||||
@@ -525,67 +539,68 @@ def test_advance_indexing_with_bool(test_varnode): | |||||
a = np.arange(9).reshape(3, 3).astype(np.float32) | a = np.arange(9).reshape(3, 3).astype(np.float32) | ||||
b = np.array([False, True, True]) | b = np.array([False, True, True]) | ||||
c = np.array([2, 0]).astype(np.int32) | c = np.array([2, 0]).astype(np.int32) | ||||
aa = Tensor(a) | |||||
bb = Tensor(b) | |||||
cc = Tensor(c) | |||||
np.testing.assert_equal(a[b, c], aa[bb, cc].numpy()) | |||||
aa = make_tensor(a, network) | |||||
bb = make_tensor(b, network) | |||||
cc = make_tensor(c, network) | |||||
np.testing.assert_equal(a[b, c], get_value(aa[bb, cc])) | |||||
a[b, c] = -1.0 | a[b, c] = -1.0 | ||||
aa[bb, cc] = -1.0 | aa[bb, cc] = -1.0 | ||||
np.testing.assert_equal(a, aa.numpy()) | |||||
np.testing.assert_equal(a, get_value(aa)) | |||||
d = np.array([-1, -2], dtype=np.float32) | d = np.array([-1, -2], dtype=np.float32) | ||||
dd = Tensor(d) | |||||
dd = make_tensor(d, network) | |||||
a[b, c] = d | a[b, c] = d | ||||
aa[bb, cc] = dd | aa[bb, cc] = dd | ||||
np.testing.assert_equal(a, aa.numpy()) | |||||
np.testing.assert_equal(a, get_value(aa)) | |||||
a = np.ones((2, 2)) | a = np.ones((2, 2)) | ||||
b = np.array([[True, False], [False, True]]) | b = np.array([[True, False], [False, True]]) | ||||
aa = Tensor(a) | |||||
bb = Tensor(b) | |||||
np.testing.assert_equal(a[b], aa[bb].numpy()) | |||||
aa = make_tensor(a, network) | |||||
bb = make_tensor(b, network) | |||||
np.testing.assert_equal(a[b], get_value(aa[bb])) | |||||
b[:] = True | b[:] = True | ||||
bb[:] = True | bb[:] = True | ||||
np.testing.assert_equal(a[b], aa[bb].numpy()) | |||||
np.testing.assert_equal(a[:, [True, False]], aa[:, [True, False]].numpy()) | |||||
np.testing.assert_equal(a[b], get_value(aa[bb])) | |||||
np.testing.assert_equal(a[:, [True, False]], get_value(aa[:, [True, False]])) | |||||
a = np.array([[True, False], [False, True]]) | a = np.array([[True, False], [False, True]]) | ||||
b = np.array([1]) | b = np.array([1]) | ||||
aa = Tensor(a) | |||||
bb = Tensor(b) | |||||
np.testing.assert_equal(a[b], aa[bb].numpy()) | |||||
aa = make_tensor(a, network) | |||||
bb = make_tensor(b, network) | |||||
np.testing.assert_equal(a[b], get_value(aa[bb])) | |||||
b = np.array([[True, True], [False, True]]) | b = np.array([[True, True], [False, True]]) | ||||
bb = Tensor(b) | |||||
np.testing.assert_equal(a[b], aa[bb].numpy()) | |||||
a[b] = False | |||||
aa[bb] = False | |||||
np.testing.assert_equal(a, aa.numpy()) | |||||
bb = make_tensor(b, network) | |||||
np.testing.assert_equal(a[b], get_value(aa[bb])) | |||||
if not test_varnode: | |||||
a[b] = False | |||||
aa[bb] = False | |||||
np.testing.assert_equal(a, get_value(aa)) | |||||
a = np.ones((2, 2), dtype=np.int32) | a = np.ones((2, 2), dtype=np.int32) | ||||
b = np.array([[False, False], [False, False]]) | b = np.array([[False, False], [False, False]]) | ||||
aa = Tensor(a) | |||||
bb = Tensor(b) | |||||
np.testing.assert_equal(a[b], aa[b].numpy()) | |||||
np.testing.assert_equal(a[b], aa[bb].numpy()) | |||||
aa = make_tensor(a, network) | |||||
bb = make_tensor(b, network) | |||||
np.testing.assert_equal(a[b], get_value(aa[b])) | |||||
np.testing.assert_equal(a[b], get_value(aa[bb])) | |||||
b = np.array([False, False]) | b = np.array([False, False]) | ||||
bb = Tensor(b) | |||||
np.testing.assert_equal(a[b], aa[bb].numpy().reshape(a[b].shape)) | |||||
bb = make_tensor(b, network) | |||||
np.testing.assert_equal(a[b], get_value(aa[bb]).reshape(a[b].shape)) | |||||
a = np.arange(576).reshape(2, 3, 4, 3, 4, 2).astype("int32") | a = np.arange(576).reshape(2, 3, 4, 3, 4, 2).astype("int32") | ||||
aa = Tensor(a) | |||||
aa = make_tensor(a, network) | |||||
b = (np.random.sample((2, 3, 4)) > 0.5).astype("bool") | b = (np.random.sample((2, 3, 4)) > 0.5).astype("bool") | ||||
bb = Tensor(b) | |||||
np.testing.assert_equal(a[b, :, 0:4:2], aa[bb, :, 0:4:2].numpy()) | |||||
bb = make_tensor(b, network) | |||||
np.testing.assert_equal(a[b, :, 0:4:2], get_value(aa[bb, :, 0:4:2])) | |||||
b = (np.random.sample((4, 3, 4)) > 0.5).astype("bool") | b = (np.random.sample((4, 3, 4)) > 0.5).astype("bool") | ||||
bb = Tensor(b) | |||||
np.testing.assert_equal(a[..., b, 0:2], aa[..., bb, 0:2].numpy()) | |||||
bb = make_tensor(b, network) | |||||
np.testing.assert_equal(a[..., b, 0:2], get_value(aa[..., bb, 0:2])) | |||||
b = (np.random.sample((3, 4, 3)) > 0.5).astype("bool") | b = (np.random.sample((3, 4, 3)) > 0.5).astype("bool") | ||||
bb = Tensor(b) | |||||
bb = make_tensor(b, network) | |||||
np.testing.assert_equal( | np.testing.assert_equal( | ||||
a[:, b, 0:2, [True, False]], aa[:, bb, 0:2, [True, False]].numpy() | |||||
a[:, b, 0:2, [True, False]], get_value(aa[:, bb, 0:2, [True, False]]) | |||||
) | ) | ||||