forked from huawei/mindspore2022
187 lines
6.4 KiB
Python
187 lines
6.4 KiB
Python
# This is the Python adaptation and derivative work of Myia (https://github.com/mila-iqia/myia/).
|
|
#
|
|
# Copyright 2020 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.
|
|
# ============================================================================
|
|
|
|
"""The names of functional part are summarized here."""
|
|
|
|
from mindspore.common._register_for_tensor import tensor_operator_registry
|
|
from .primitive import Primitive
|
|
from . import operations as P
|
|
from .operations import _grad_ops
|
|
from .._extends import builtin_operations as BP
|
|
|
|
typeof = Primitive('typeof')
|
|
hastype = Primitive('hastype')
|
|
cast = P.Cast()
|
|
dtype = P.DType()
|
|
isconstant = Primitive('is_constant')
|
|
isconstant.add_prim_attr('const_value', True)
|
|
|
|
|
|
issubclass_ = P.IsSubClass()
|
|
isinstance_ = P.IsInstance()
|
|
fill = P.Fill()
|
|
tile = P.Tile()
|
|
select = P.Select()
|
|
size = P.Size()
|
|
ones_like = P.OnesLike()
|
|
shape = P.Shape()
|
|
rank = P.Rank()
|
|
reshape = P.Reshape()
|
|
# control_depend: represent dependency between two operators
|
|
control_depend = P.ControlDepend()
|
|
merge = P.Merge()
|
|
geswitch = P.GeSwitch()
|
|
addn = P.AddN()
|
|
tensor_add = P.TensorAdd()
|
|
neg_tensor = P.Neg()
|
|
tensor_lt = P.Less()
|
|
tensor_le = P.LessEqual()
|
|
tensor_gt = P.Greater()
|
|
tensor_ge = P.GreaterEqual()
|
|
tensor_sub = P.Sub()
|
|
tensor_mul = P.Mul()
|
|
tensor_div = P.RealDiv()
|
|
tensor_floordiv = P.FloorDiv()
|
|
tensor_pow = P.Pow()
|
|
tensor_mod = P.FloorMod()
|
|
strided_slice = P.StridedSlice()
|
|
same_type_shape = P.SameTypeShape()
|
|
check_bprop = P.CheckBprop()
|
|
equal = P.Equal()
|
|
not_equal = P.NotEqual()
|
|
assign_sub = P.AssignSub()
|
|
assign = P.Assign()
|
|
square = P.Square()
|
|
sqrt = P.Sqrt()
|
|
scalar_to_array = P.ScalarToArray()
|
|
scalar_to_tensor = P.ScalarToTensor()
|
|
tuple_to_array = P.TupleToArray()
|
|
scalar_cast = P.ScalarCast()
|
|
print_ = P.Print()
|
|
expand_dims = P.ExpandDims()
|
|
scatter_nd = P.ScatterNd()
|
|
gather = P.GatherV2()
|
|
gather_nd = P.GatherNd()
|
|
scatter_update = P.ScatterUpdate()
|
|
scatter_nd_update = P.ScatterNdUpdate()
|
|
pack = P.Pack()
|
|
partial = P.Partial()
|
|
# depend: mount a node to another node
|
|
depend = P.Depend()
|
|
|
|
|
|
tuple_setitem = Primitive('tuple_setitem')
|
|
tuple_getitem = Primitive('tuple_getitem')
|
|
list_getitem = Primitive('list_getitem')
|
|
list_setitem = Primitive('list_setitem')
|
|
dict_getitem = Primitive('dict_getitem')
|
|
dict_setitem = Primitive('dict_setitem')
|
|
tuple_div = Primitive("tuple_div")
|
|
tuple_len = Primitive("tuple_len")
|
|
tuple_reversed = Primitive("tuple_reversed")
|
|
make_range = Primitive("make_range")
|
|
make_tuple = Primitive('make_tuple')
|
|
make_dict = Primitive('make_dict')
|
|
make_list = Primitive('make_list')
|
|
make_slice = Primitive('make_slice')
|
|
tuple_equal = Primitive("tuple_equal")
|
|
list_equal = Primitive("list_equal")
|
|
make_ref = Primitive("make_ref")
|
|
|
|
|
|
scalar_add = Primitive('scalar_add')
|
|
scalar_mul = Primitive('scalar_mul')
|
|
scalar_sub = Primitive('scalar_sub')
|
|
scalar_div = Primitive('scalar_div')
|
|
scalar_floordiv = Primitive('scalar_floordiv')
|
|
scalar_log = Primitive('scalar_log')
|
|
scalar_pow = Primitive('scalar_pow')
|
|
scalar_gt = Primitive('scalar_gt')
|
|
scalar_ge = Primitive('scalar_ge')
|
|
scalar_le = Primitive('scalar_le')
|
|
scalar_lt = Primitive('scalar_lt')
|
|
scalar_eq = Primitive('scalar_eq')
|
|
scalar_ne = Primitive('scalar_ne')
|
|
scalar_uadd = Primitive('scalar_uadd')
|
|
scalar_usub = Primitive('scalar_usub')
|
|
scalar_mod = Primitive('scalar_mod')
|
|
string_eq = Primitive('string_equal')
|
|
string_concat = Primitive('string_concat')
|
|
bool_not = Primitive("bool_not")
|
|
bool_or = Primitive("bool_or")
|
|
bool_and = Primitive("bool_and")
|
|
logical_and = P.LogicalAnd()
|
|
logical_or = P.LogicalOr()
|
|
logical_not = P.LogicalNot()
|
|
array_to_scalar = Primitive('array_to_scalar')
|
|
is_ = Primitive("is_")
|
|
is_not = Primitive("is_not")
|
|
in_dict = Primitive("in_dict")
|
|
not_in_dict = Primitive("not_in_dict")
|
|
mixed_precision_cast = Primitive("mixed_precision_cast")
|
|
broadcast_gradient_args = Primitive('BroadcastGradientArgs')
|
|
dot = Primitive('dot')
|
|
array_reduce = Primitive('array_reduce')
|
|
zeros_like = P.ZerosLike()
|
|
identity = Primitive('identity')
|
|
distribute = Primitive('distribute')
|
|
embed = Primitive('embed')
|
|
ref_to_embed = _grad_ops.RefToEmbed()
|
|
env_setitem = Primitive('env_setitem')
|
|
env_getitem = Primitive('env_getitem')
|
|
env_add = Primitive('env_add')
|
|
J = Primitive('J')
|
|
switch = Primitive('switch')
|
|
switch_layer = Primitive('switch_layer')
|
|
# for sum bprop
|
|
reduced_shape = Primitive("reduced_shape")
|
|
# shape_mul:input mush be shape multiply elemts in tuple(shape)
|
|
shape_mul = Primitive("shape_mul")
|
|
# a primitive to compare between tuple.
|
|
stop_gradient = Primitive("stop_gradient")
|
|
|
|
|
|
make_indexed_slices = Primitive('MakeIndexedSlices')
|
|
indexed_slices_get_values = Primitive('IndexedSlicesGetValues')
|
|
indexed_slices_get_indices = Primitive('IndexedSlicesGetIndices')
|
|
indexed_slices_get_dense_shape = Primitive('IndexedSlicesGetDenseShape')
|
|
|
|
make_sparse_tensor = Primitive('MakeSparseTensor')
|
|
sparse_tensor_get_values = Primitive('SparseTensorGetValues')
|
|
sparse_tensor_get_indices = Primitive('SparseTensorGetIndices')
|
|
sparse_tensor_get_dense_shape = Primitive('SparseTensorGetDenseShape')
|
|
|
|
tensor_operator_registry.register('__add__', tensor_add)
|
|
tensor_operator_registry.register('__sub__', tensor_sub)
|
|
tensor_operator_registry.register('__mul__', tensor_mul)
|
|
tensor_operator_registry.register('__truediv__', tensor_div)
|
|
tensor_operator_registry.register('__mod__', tensor_mod)
|
|
tensor_operator_registry.register('__pow__', tensor_pow)
|
|
tensor_operator_registry.register('__floordiv__', tensor_floordiv)
|
|
#ms cannot support Tensor(True) compare
|
|
tensor_operator_registry.register('__eq__', equal)
|
|
tensor_operator_registry.register('__ne__', not_equal)
|
|
tensor_operator_registry.register('__neg__', neg_tensor)
|
|
tensor_operator_registry.register('__lt__', tensor_lt)
|
|
tensor_operator_registry.register('__le__', tensor_le)
|
|
tensor_operator_registry.register('__gt__', tensor_gt)
|
|
tensor_operator_registry.register('__ge__', tensor_ge)
|
|
tensor_operator_registry.register('shape', shape)
|
|
#support GE backend for no compare operators
|
|
tensor_operator_registry.register('vm_compare', BP.vm_compare)
|
|
tensor_operator_registry.register('cast', cast)
|