608 lines
19 KiB
Python
608 lines
19 KiB
Python
# Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# 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.
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import typing
|
|
from collections import defaultdict
|
|
from collections.abc import Callable, Sequence
|
|
from typing import Any, TypeGuard, TypeVar
|
|
from uuid import uuid4
|
|
from weakref import WeakKeyDictionary
|
|
|
|
import numpy as np
|
|
|
|
import paddle
|
|
from paddle.pir.core import convert_nptype_to_datatype
|
|
|
|
from ..base.data_feeder import check_dtype, convert_dtype
|
|
from ..base.framework import (
|
|
Block,
|
|
Variable,
|
|
in_dygraph_mode,
|
|
)
|
|
from ..pir import Value
|
|
|
|
if typing.TYPE_CHECKING:
|
|
from paddle._typing import NestedStructure, ShapeLike
|
|
|
|
_T = TypeVar("_T")
|
|
_U = TypeVar("_U")
|
|
|
|
|
|
class NotSupportedTensorArgumentError(TypeError):
|
|
def __init__(self, msg, name: str):
|
|
super().__init__(msg)
|
|
self.name = name
|
|
|
|
|
|
def convert_to_list(value, n, name, dtype=int):
|
|
"""
|
|
Converts a single numerical type or iterable of numerical
|
|
types into a numerical type list.
|
|
|
|
Arguments:
|
|
value: The value to validate and convert. Could an int, or any iterable
|
|
of ints.
|
|
n: The size of the list to be returned.
|
|
name: The name of the argument being validated, e.g. "stride" or
|
|
"filter_size". This is only used to format error messages.
|
|
dtype: the numerical type of the element of the list to be returned.
|
|
|
|
Returns:
|
|
A list of n dtypes.
|
|
|
|
Raises:
|
|
ValueError: If something else than an int/long or iterable thereof was
|
|
passed.
|
|
"""
|
|
if isinstance(value, dtype):
|
|
return [value] * n
|
|
else:
|
|
if isinstance(value, (Variable, paddle.pir.Value)):
|
|
raise NotSupportedTensorArgumentError(
|
|
f"`{name}` required numerical type with `{dtype}`, but received Tensor.",
|
|
name,
|
|
)
|
|
try:
|
|
value_list = list(value)
|
|
except TypeError:
|
|
raise ValueError(
|
|
f"The {name}'s type must be list or tuple. Received: {value}"
|
|
)
|
|
if len(value_list) != n:
|
|
raise ValueError(
|
|
f"The {name}'s length must be {n}. Received: {value}"
|
|
)
|
|
for single_value in value_list:
|
|
if isinstance(single_value, (Variable, paddle.pir.Value)):
|
|
raise NotSupportedTensorArgumentError(
|
|
f"`{name}` required numerical type with `{dtype}`, but received Tensor.",
|
|
name,
|
|
)
|
|
try:
|
|
dtype(single_value)
|
|
except (ValueError, TypeError):
|
|
raise ValueError(
|
|
f"The {name}'s type must be a list or tuple of {n} {dtype}. "
|
|
+ f"Received: {value} including element {single_value} of type {type(single_value)}"
|
|
)
|
|
return value_list
|
|
|
|
|
|
def is_sequence(seq: Any) -> TypeGuard[typing.Sequence[Any] | dict[str, Any]]:
|
|
"""
|
|
Whether `seq` is an entry or nested structure
|
|
"""
|
|
if isinstance(seq, dict):
|
|
return True
|
|
return isinstance(seq, Sequence) and not isinstance(seq, str)
|
|
|
|
|
|
class UniqueIdMap(WeakKeyDictionary):
|
|
def __init__(self):
|
|
super().__init__(self)
|
|
self.data = defaultdict(uuid4)
|
|
|
|
|
|
uniqueidmap = UniqueIdMap()
|
|
|
|
|
|
def uniqueid(obj):
|
|
if isinstance(obj, str):
|
|
return (hash(obj),)
|
|
elif isinstance(obj, list):
|
|
return (id(obj),)
|
|
else:
|
|
return (uniqueidmap[obj].int,)
|
|
|
|
|
|
def _hash_with_id(*args):
|
|
"""
|
|
Return int hash value calculated by id(arg) or tuple(id1,id2, ...).
|
|
"""
|
|
assert len(args) > 0
|
|
info = ()
|
|
for v in args:
|
|
info = info + uniqueid(v)
|
|
return hash(info)
|
|
|
|
|
|
def _sorted(dict_):
|
|
"""
|
|
Returns a sorted list of the dict keys, with error if keys not sortable.
|
|
"""
|
|
try:
|
|
return sorted(dict_.keys())
|
|
except TypeError:
|
|
raise TypeError("nest only supports dicts with sortable keys.")
|
|
|
|
|
|
def _yield_value(iterable):
|
|
if isinstance(iterable, dict):
|
|
for key in _sorted(iterable):
|
|
yield iterable[key]
|
|
else:
|
|
yield from iterable
|
|
|
|
|
|
def _yield_flat_nest(nest):
|
|
for n in _yield_value(nest):
|
|
if is_sequence(n):
|
|
yield from _yield_flat_nest(n)
|
|
else:
|
|
yield n
|
|
|
|
|
|
def to_sequence(nest):
|
|
if is_sequence(nest):
|
|
return nest
|
|
else:
|
|
return [nest]
|
|
|
|
|
|
def flatten(nest: NestedStructure[_T]) -> typing.Sequence[_T]:
|
|
"""
|
|
:alias_main: paddle.flatten
|
|
:alias: paddle.flatten,paddle.tensor.flatten,paddle.tensor.manipulation.flatten
|
|
:old_api: paddle.base.layers.flatten
|
|
|
|
Traverse all entries in the nested structure and put them into an list.
|
|
"""
|
|
if is_sequence(nest):
|
|
return list(_yield_flat_nest(nest))
|
|
else:
|
|
return [nest]
|
|
|
|
|
|
def _sequence_like(instance, args):
|
|
"""
|
|
Convert the sequence `args` to the same type as `instance`.
|
|
"""
|
|
if isinstance(instance, dict):
|
|
result = dict(zip(_sorted(instance), args))
|
|
return type(instance)((key, result[key]) for key in instance.keys())
|
|
elif (
|
|
isinstance(instance, tuple)
|
|
and hasattr(instance, "_fields")
|
|
and isinstance(instance._fields, Sequence)
|
|
and all(isinstance(f, str) for f in instance._fields)
|
|
):
|
|
# This is a namedtuple
|
|
return type(instance)(*args)
|
|
else:
|
|
# Not a namedtuple
|
|
return type(instance)(args)
|
|
|
|
|
|
def _packed_nest_with_indices(structure, flat, index):
|
|
"""
|
|
Helper function for pack_sequence_as.
|
|
"""
|
|
packed = []
|
|
for s in _yield_value(structure):
|
|
if is_sequence(s):
|
|
new_index, child = _packed_nest_with_indices(s, flat, index)
|
|
packed.append(_sequence_like(s, child))
|
|
index = new_index
|
|
else:
|
|
# Paddle requires python version > 3.7, so dict is always OrderedDict
|
|
packed.append(
|
|
flat[index]
|
|
if not isinstance(flat, dict)
|
|
else list(flat.values())[index]
|
|
)
|
|
index += 1
|
|
return index, packed
|
|
|
|
|
|
def pack_sequence_as(structure, flat_sequence):
|
|
"""
|
|
Pack a given flattened sequence into a given structure.
|
|
"""
|
|
if not is_sequence(flat_sequence):
|
|
raise TypeError("flat_sequence must be a sequence")
|
|
if not is_sequence(structure):
|
|
if len(flat_sequence) != 1:
|
|
raise ValueError(
|
|
f"Structure is a scalar but len(flat_sequence) == {len(flat_sequence)} > 1"
|
|
)
|
|
return flat_sequence[0]
|
|
flat_structure = flatten(structure)
|
|
if len(flat_structure) != len(flat_sequence):
|
|
raise ValueError(
|
|
f"Could not pack sequence. Structure had {len(flat_structure)} elements, but flat_sequence "
|
|
f"had {len(flat_sequence)} elements. Structure: {structure}, flat_sequence: {flat_sequence}."
|
|
)
|
|
_, packed = _packed_nest_with_indices(structure, flat_sequence, 0)
|
|
return _sequence_like(structure, packed)
|
|
|
|
|
|
def map_structure(
|
|
func: Callable[[_T], _U], *structure: NestedStructure[_T]
|
|
) -> NestedStructure[_U]:
|
|
"""
|
|
Apply `func` to each entry in `structure` and return a new structure.
|
|
"""
|
|
flat_structure = [flatten(s) for s in structure]
|
|
entries = zip(*flat_structure)
|
|
return pack_sequence_as(structure[0], [func(*x) for x in entries])
|
|
|
|
|
|
def hold_mutable_vars(structure):
|
|
"""
|
|
Returns whether structure holds sequence like `list/dict`.
|
|
"""
|
|
for s in structure:
|
|
if is_sequence(s):
|
|
return True
|
|
return False
|
|
|
|
|
|
def copy_mutable_vars(structure):
|
|
"""
|
|
Returns vars copied from sequence without mutable property.
|
|
"""
|
|
flat_structure = copy.copy(flatten(structure))
|
|
return pack_sequence_as(structure, flat_structure)
|
|
|
|
|
|
def _recursive_assert_same_structure(nest1, nest2, check_types, skip_if):
|
|
"""
|
|
Helper function for `assert_same_structure`.
|
|
"""
|
|
if skip_if is not None and (skip_if(nest1) or skip_if(nest2)):
|
|
return
|
|
is_sequence_nest1 = is_sequence(nest1)
|
|
if is_sequence_nest1 != is_sequence(nest2):
|
|
raise ValueError(
|
|
"The two structures don't have the same nested structure.\n\n"
|
|
f"First structure: {nest1}\n\nSecond structure: {nest2}."
|
|
)
|
|
if not is_sequence_nest1:
|
|
return # finished checking
|
|
if check_types:
|
|
type_nest1 = type(nest1)
|
|
type_nest2 = type(nest2)
|
|
if type_nest1 != type_nest2:
|
|
raise TypeError(
|
|
"The two structures don't have the same sequence type. First "
|
|
f"structure has type {type_nest1}, while second structure has type {type_nest2}."
|
|
)
|
|
if isinstance(nest1, dict):
|
|
keys1 = set(nest1.keys())
|
|
keys2 = set(nest2.keys())
|
|
if keys1 != keys2:
|
|
raise ValueError(
|
|
"The two dictionaries don't have the same set of keys. First "
|
|
f"structure has keys {keys1}, while second structure has keys {keys2}."
|
|
)
|
|
nest1_as_sequence = list(_yield_value(nest1))
|
|
nest2_as_sequence = list(_yield_value(nest2))
|
|
if len(nest1_as_sequence) != len(nest2_as_sequence):
|
|
raise ValueError(
|
|
"The two structures don't have the same number of elements.\n\n"
|
|
f"First structure ({len(nest1_as_sequence)} elements): {nest1}\n\n"
|
|
f"Second structure ({len(nest2_as_sequence)} elements): {nest2}"
|
|
)
|
|
for n1, n2 in zip(nest1_as_sequence, nest2_as_sequence):
|
|
_recursive_assert_same_structure(n1, n2, check_types, skip_if)
|
|
|
|
|
|
def padding_to_same_structure(nest1, nest2, obj=None):
|
|
def _padding_to_same_structure_single(value, obj):
|
|
def change_none_to_obj(x):
|
|
if x is None:
|
|
return obj
|
|
return x
|
|
|
|
if is_sequence(value):
|
|
value = pack_sequence_as(
|
|
value, [change_none_to_obj(item) for item in flatten(value)]
|
|
)
|
|
else:
|
|
value = change_none_to_obj(value)
|
|
return value
|
|
|
|
nest1 = _padding_to_same_structure_single(nest1, obj)
|
|
nest2 = _padding_to_same_structure_single(nest2, obj)
|
|
return nest1, nest2
|
|
|
|
|
|
def assert_same_structure(nest1, nest2, check_types=True, skip_if=None):
|
|
"""
|
|
Confirm two nested structures with the same structure.
|
|
"""
|
|
if skip_if is not None and (skip_if(nest1) or skip_if(nest2)):
|
|
return
|
|
len_nest1 = len(flatten(nest1)) if is_sequence(nest1) else 1
|
|
len_nest2 = len(flatten(nest2)) if is_sequence(nest2) else 1
|
|
if len_nest1 != len_nest2 and skip_if is None:
|
|
raise ValueError(
|
|
"The two structures don't have the same number of "
|
|
f"elements.\n\nFirst structure ({len_nest1} elements): {nest1}\n\n"
|
|
f"Second structure ({len_nest2} elements): {nest2}"
|
|
)
|
|
_recursive_assert_same_structure(nest1, nest2, check_types, skip_if)
|
|
|
|
|
|
def _is_symmetric_padding(padding, data_dim):
|
|
"""
|
|
Check whether padding is symmetrical.
|
|
"""
|
|
assert len(padding) == data_dim * 2 or len(padding) == data_dim
|
|
is_sys = True
|
|
if len(padding) == data_dim * 2:
|
|
for i in range(data_dim):
|
|
if padding[i * 2] != padding[i * 2 + 1]:
|
|
is_sys = False
|
|
return is_sys
|
|
|
|
|
|
def _contain_var(list_or_tuple):
|
|
"""
|
|
Check whether list or tuple contains variable / Value.
|
|
"""
|
|
for item in list_or_tuple:
|
|
if isinstance(item, (Variable, paddle.pir.Value)):
|
|
return True
|
|
return False
|
|
|
|
|
|
def get_int_tensor_list(ele_list, default_dtype='int64'):
|
|
int_tensor_list = []
|
|
for ele in ele_list:
|
|
if isinstance(ele, paddle.pir.Value):
|
|
ele.stop_gradient = True
|
|
if convert_dtype(ele.dtype) != default_dtype:
|
|
ele = paddle.cast(x=ele, dtype=default_dtype)
|
|
if ele.shape != []:
|
|
ele = paddle.reshape(ele, [])
|
|
int_tensor_list.append(ele)
|
|
else:
|
|
temp_out = paddle.tensor.fill_constant(
|
|
shape=[],
|
|
dtype=convert_nptype_to_datatype(np.dtype(default_dtype)),
|
|
value=ele,
|
|
force_cpu=True,
|
|
)
|
|
int_tensor_list.append(temp_out)
|
|
return int_tensor_list
|
|
|
|
|
|
def get_shape_tensor_inputs(inputs, attrs, shape, op_type):
|
|
from paddle.tensor import fill_constant
|
|
|
|
def _get_attr_shape(list_shape):
|
|
attr_shape = []
|
|
for idx, dim in enumerate(list_shape):
|
|
if isinstance(dim, Variable):
|
|
attr_shape.append(-1)
|
|
else:
|
|
attr_shape.append(dim)
|
|
return attr_shape
|
|
|
|
def _get_shape_tensor(list_shape):
|
|
shape_tensor_list = []
|
|
for idx, dim in enumerate(list_shape):
|
|
if isinstance(dim, Variable):
|
|
dim.stop_gradient = True
|
|
check_dtype(
|
|
dim.dtype,
|
|
'shape[' + str(idx) + ']',
|
|
['int32', 'int64'],
|
|
op_type,
|
|
f'(When type of shape in {op_type} is list or tuple.)',
|
|
)
|
|
if convert_dtype(dim.dtype) == 'int64':
|
|
dim = paddle.cast(x=dim, dtype='int32')
|
|
shape_tensor_list.append(dim)
|
|
else:
|
|
temp_out = fill_constant([], 'int32', dim, force_cpu=True)
|
|
shape_tensor_list.append(temp_out)
|
|
return shape_tensor_list
|
|
|
|
if isinstance(shape, Variable):
|
|
shape.stop_gradient = True
|
|
check_dtype(
|
|
shape.dtype,
|
|
'shape',
|
|
['int32', 'int64'],
|
|
'fill_constant',
|
|
f'(When type of shape in {op_type} is Variable.)',
|
|
)
|
|
if convert_dtype(shape.dtype) == 'int64':
|
|
shape = paddle.cast(shape, 'int32')
|
|
inputs["ShapeTensor"] = shape
|
|
elif isinstance(shape, (list, tuple)):
|
|
attrs["shape"] = _get_attr_shape(shape)
|
|
if _contain_var(shape):
|
|
inputs['ShapeTensorList'] = _get_shape_tensor(shape)
|
|
else:
|
|
raise TypeError("Shape only supports Variable, or list, or tuple.")
|
|
|
|
|
|
def _convert_to_tensor_list(old_list, dtype="int32"):
|
|
"""
|
|
Converts all elements of a list to Variable / Value.
|
|
"""
|
|
from paddle.tensor import fill_constant
|
|
|
|
if _contain_var(old_list):
|
|
for ele in old_list:
|
|
if isinstance(ele, paddle.pir.Value):
|
|
dtype = ele.dtype
|
|
|
|
new_list_tensor = []
|
|
for ele in old_list:
|
|
if isinstance(ele, (Variable, paddle.pir.Value)):
|
|
ele.stop_gradient = True
|
|
new_list_tensor.append(ele)
|
|
else:
|
|
assert isinstance(ele, int)
|
|
temp_out = fill_constant([1], dtype, ele, force_cpu=True)
|
|
new_list_tensor.append(temp_out)
|
|
return new_list_tensor
|
|
|
|
|
|
def convert_shape_to_list(shape):
|
|
"""
|
|
Convert shape(list, tuple, variable) to list in imperative mode
|
|
"""
|
|
if isinstance(shape, (list, tuple)):
|
|
shape_out = []
|
|
for x in shape:
|
|
if isinstance(x, Variable):
|
|
# skip item if size = 0
|
|
if x.size > 0:
|
|
shape_out.append(x.item(0))
|
|
else:
|
|
shape_out.append(x)
|
|
shape = shape_out
|
|
else:
|
|
if in_dygraph_mode():
|
|
shape = shape.astype(int).tolist()
|
|
return shape
|
|
|
|
|
|
def check_shape(shape):
|
|
"""
|
|
Check shape type and shape elements type before passing it to fill_constant
|
|
"""
|
|
if isinstance(shape, (Variable, Value)):
|
|
check_dtype(shape.dtype, 'shape', ['int32', 'int64'], 'fill_constant')
|
|
elif isinstance(shape, (list, tuple)):
|
|
for ele in shape:
|
|
if not isinstance(ele, (Variable, Value)):
|
|
if ele < 0:
|
|
raise ValueError(
|
|
"All elements in ``shape`` must be positive when it's a list or tuple"
|
|
)
|
|
if not isinstance(ele, int):
|
|
raise TypeError(
|
|
"All elements in ``shape`` must be integers when it's a list or tuple"
|
|
)
|
|
else:
|
|
check_dtype(
|
|
ele.dtype,
|
|
'element of shape',
|
|
['int32', 'int64'],
|
|
'fill_constant',
|
|
)
|
|
|
|
|
|
def try_get_constant_shape_from_tensor(shape_tensor):
|
|
"""Try to get shape from a tensor with constant value.
|
|
|
|
For example,
|
|
|
|
import paddle
|
|
paddle.enable_static()
|
|
data = paddle.static.data(name="x", shape=[-1, 2], dtype='float32')
|
|
shape = paddle.shape(data) # shape should be [-1, 2] instead of [-1, -1]
|
|
x = paddle.uniform(shape)
|
|
print(x.shape)
|
|
# (-1, 2)
|
|
|
|
"""
|
|
if not in_dygraph_mode():
|
|
try:
|
|
if shape_tensor.op is not None:
|
|
generate_op = shape_tensor.op
|
|
if generate_op.type == 'shape':
|
|
var = shape_tensor.block.vars[
|
|
generate_op.input_arg_names[0]
|
|
]
|
|
return var.shape
|
|
except:
|
|
return None
|
|
|
|
return None
|
|
|
|
|
|
def get_inputs_outputs_in_block(block):
|
|
"""
|
|
Returns the inputs and outputs variable used in this block but not
|
|
created in this block.
|
|
"""
|
|
assert isinstance(block, Block), (
|
|
"input non-Block argument for get_inputs_outputs_in_block."
|
|
)
|
|
assert block.parent_idx != -1, (
|
|
"input block should be a sub-block, not main block."
|
|
)
|
|
|
|
# Find input/output var names of all ops in block
|
|
inner_inputs = set()
|
|
inner_outputs = set()
|
|
for op in block.ops:
|
|
for iname in op.input_names:
|
|
for in_var_name in op.input(iname):
|
|
if not block.has_var(in_var_name):
|
|
# variable not created in this block
|
|
inner_inputs.add(in_var_name)
|
|
for oname in op.output_names:
|
|
for out_var_name in op.output(oname):
|
|
if not block.has_var(out_var_name):
|
|
# variable not created in this block
|
|
inner_outputs.add(out_var_name)
|
|
|
|
return inner_inputs, inner_outputs
|
|
|
|
|
|
def is_same_shape(shape1: ShapeLike, shape2: ShapeLike) -> bool:
|
|
"""
|
|
Check whether two shapes are the same. Deal with the dynamic shape.
|
|
"""
|
|
if paddle.in_dynamic_mode():
|
|
return shape1 == shape2
|
|
|
|
def is_tensor(x):
|
|
return isinstance(x, (paddle.static.Variable, paddle.pir.Value))
|
|
|
|
def is_dynamic_axis(axis):
|
|
return is_tensor(axis) or axis == -1
|
|
|
|
if is_tensor(shape1) or is_tensor(shape2):
|
|
return True
|
|
if len(shape1) != len(shape2):
|
|
return False
|
|
for s1, s2 in zip(shape1, shape2):
|
|
if is_dynamic_axis(s1) or is_dynamic_axis(s2):
|
|
continue
|
|
if s1 != s2:
|
|
return False
|
|
return True
|