chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,878 @@
|
||||
# Copyright (c) 2022 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 re
|
||||
import warnings
|
||||
from contextlib import contextmanager
|
||||
|
||||
import paddle
|
||||
from paddle.autograd.py_layer import PyLayerMeta
|
||||
from paddle.base.data_feeder import convert_dtype
|
||||
from paddle.base.dygraph.base import _convert_into_variable, in_to_static_mode
|
||||
from paddle.base.framework import Variable, core, default_main_program
|
||||
from paddle.framework import use_pir_api
|
||||
from paddle.jit.utils import OrderedSet
|
||||
from paddle.pir import Value
|
||||
from paddle.static.amp.fp16_utils import AmpOptions
|
||||
from paddle.utils import is_sequence, map_structure
|
||||
|
||||
from .py_layer import StaticPyLayer
|
||||
from .utils import (
|
||||
RETURN_NO_VALUE_VAR_NAME,
|
||||
Dygraph2StaticException,
|
||||
GetterSetterHelper,
|
||||
UndefinedVar,
|
||||
create_undefined_variable,
|
||||
)
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
def to_static_variable(x, dtype=None):
|
||||
'''
|
||||
Translate a Python Tensor to PaddlePaddle static graph Tensor
|
||||
'''
|
||||
if isinstance(x, bool):
|
||||
dtype = 'bool' if dtype is None else dtype
|
||||
return paddle.full(shape=[], dtype=dtype, fill_value=x)
|
||||
if isinstance(x, float):
|
||||
dtype = 'float64' if dtype is None else dtype
|
||||
return paddle.full(shape=[], dtype=dtype, fill_value=x)
|
||||
if isinstance(x, int):
|
||||
dtype = 'int64' if dtype is None else dtype
|
||||
return paddle.full(shape=[], dtype=dtype, fill_value=x)
|
||||
if not use_pir_api() and (isinstance(x, UndefinedVar) or x is None):
|
||||
"""
|
||||
for early return case, we need a variable to represent None, current we use data_layer_not_check.
|
||||
"""
|
||||
return create_undefined_variable()
|
||||
if is_sequence(x):
|
||||
return map_structure(to_static_variable, x)
|
||||
return x
|
||||
|
||||
|
||||
def convert_attr(x, attr):
|
||||
# TODO(cleanup-legacy-ir): In PIR mode, the size attr in
|
||||
# Value and Tensor are unified. So we don't need to transform
|
||||
# the size attr into a method call. The AttributeJstTransformer and
|
||||
# convert_attr can be safely removed.
|
||||
if (
|
||||
isinstance(x, Variable)
|
||||
and not isinstance(x, paddle.Tensor)
|
||||
and attr == "size"
|
||||
):
|
||||
return x.size()
|
||||
else:
|
||||
return getattr(x, attr)
|
||||
|
||||
|
||||
def convert_load(x):
|
||||
# convert dygraph `PyLayer` into StaticPyLayer
|
||||
if isinstance(x, PyLayerMeta):
|
||||
return StaticPyLayer(x)
|
||||
|
||||
if in_to_static_mode():
|
||||
if isinstance(x, paddle.Tensor):
|
||||
"""
|
||||
TODO:(@xiongkun) may run convert_load in dygraph mode, which should be fixed.
|
||||
"""
|
||||
return _convert_into_variable(x)
|
||||
|
||||
# get the new output of the var
|
||||
if isinstance(x, Value):
|
||||
from paddle.jit.dy2static.parameter_recorder import (
|
||||
_global_inplace_map,
|
||||
)
|
||||
|
||||
new_var = _global_inplace_map.get(
|
||||
paddle.static.default_main_program(), x
|
||||
)
|
||||
if new_var is not None:
|
||||
return new_var
|
||||
|
||||
if isinstance(x, Variable):
|
||||
cur_block = default_main_program().current_block()
|
||||
|
||||
from paddle.jit.dy2static.program_translator import (
|
||||
ProgramTranslator,
|
||||
)
|
||||
|
||||
new_var = ProgramTranslator.get_instance()._inplace_map.get(
|
||||
cur_block.program, x.desc.id()
|
||||
)
|
||||
if new_var is not None:
|
||||
return new_var
|
||||
|
||||
if x is paddle.amp.auto_cast and not use_pir_api():
|
||||
return convert_auto_cast
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def indexable(x, code=None):
|
||||
if isinstance(x, (Variable, Value)):
|
||||
return x
|
||||
elif hasattr(x, '__iter__'):
|
||||
return list(x)
|
||||
elif hasattr(x, '__len__') and hasattr(
|
||||
x, '__getitem__'
|
||||
): # used for customed type and non-iterable type.
|
||||
return x
|
||||
else:
|
||||
raise RuntimeError("X can't be convert into indexable.")
|
||||
|
||||
|
||||
def unpack_by_structure(target, structure):
|
||||
"""unified unpack interface for paddle and python."""
|
||||
if isinstance(target, (Variable, Value)):
|
||||
return _unpack_by_structure_paddle(target, structure)
|
||||
else:
|
||||
return _unpack_by_structure_python(target, structure)
|
||||
|
||||
|
||||
def _unpack_by_structure_python(target, structure):
|
||||
"""TODO(xiongkun): analysis the differences between python and paddle unpack."""
|
||||
return _unpack_by_structure_paddle(target, structure)
|
||||
|
||||
|
||||
def _unpack_by_structure_paddle(target, structure):
|
||||
if structure == 1:
|
||||
return target
|
||||
ret = []
|
||||
for idx, ele in enumerate(structure):
|
||||
if ele == 1:
|
||||
ret.append(target[idx])
|
||||
continue
|
||||
if isinstance(ele, list):
|
||||
ret.append(unpack_by_structure(target[idx], ele))
|
||||
continue
|
||||
raise AssertionError("structure element must be 1 or list")
|
||||
return ret
|
||||
|
||||
|
||||
def convert_while_loop(
|
||||
cond, body, getter, setter, return_name_ids=None, push_pop_names=None
|
||||
):
|
||||
"""
|
||||
A function representation of a Python ``while`` statement.
|
||||
|
||||
Args:
|
||||
cond(Callable): A callable object that returns a boolean variable to control whether to execute the loop body. It takes ``loop_vars`` as arguments.
|
||||
body(Callable): A callable object that returns a tuple or list of variables with the same arguments ``loops_vars`` as ``cond`` .
|
||||
get_args(callable): Get all arguments that needed in true_fn and false_fn.
|
||||
set_args(callable): Update arguments that modified in trure_fn and false_fn.
|
||||
return_name_ids(list[string], optional): the returned names.
|
||||
push_pop_names(list[string], optional): the names on which called .append() or .pop().
|
||||
|
||||
Returns:
|
||||
A list or tuple of variables which returned by ``body``.
|
||||
"""
|
||||
|
||||
# NOTE: It may be slower if cond is very expensive, but usually cond is just O(1).
|
||||
# If loop_vars is changed during cond callable, then it causes bug, but current logical_and/logical_not/... doesn't change the loop_vars.
|
||||
pred = cond()
|
||||
if isinstance(pred, (Variable, Value)):
|
||||
_run_paddle_while(
|
||||
cond, body, getter, setter, return_name_ids, push_pop_names
|
||||
)
|
||||
else:
|
||||
_run_py_while(cond, body, getter, setter)
|
||||
|
||||
|
||||
def _convert_tensor_array_if_necessary(setterhelper, push_pop_names):
|
||||
push_pop_vars = setterhelper.get(push_pop_names)
|
||||
if push_pop_vars is None:
|
||||
return
|
||||
|
||||
def maybe_to_tensor_array(v):
|
||||
if isinstance(v, list):
|
||||
dtype = (
|
||||
paddle.base.libpaddle.DataType.UNDEFINED
|
||||
if use_pir_api()
|
||||
else "float32"
|
||||
)
|
||||
return paddle.tensor.create_array(dtype, initialized_list=v)
|
||||
else:
|
||||
return v
|
||||
|
||||
setterhelper.set(
|
||||
push_pop_names, [maybe_to_tensor_array(v) for v in push_pop_vars]
|
||||
)
|
||||
|
||||
|
||||
def _run_paddle_while(
|
||||
cond, body, getter, setter, return_name_ids, push_pop_names
|
||||
):
|
||||
# NOTE: loop_vars of Paddle op `control_flow.while_loop` must be Paddle Tensors.
|
||||
helper = GetterSetterHelper(getter, setter, return_name_ids, push_pop_names)
|
||||
_convert_tensor_array_if_necessary(helper, push_pop_names)
|
||||
|
||||
union_name = (
|
||||
OrderedSet(return_name_ids) if return_name_ids else OrderedSet()
|
||||
) | (OrderedSet(push_pop_names) if push_pop_names else OrderedSet())
|
||||
union_name = list(union_name)
|
||||
|
||||
def new_body_fn(*args):
|
||||
"""wrap the body() and add return value for `while_loop`
|
||||
the args may be differ from getter().
|
||||
"""
|
||||
mutable_loop_vars = args
|
||||
helper.set(union_name, mutable_loop_vars)
|
||||
body()
|
||||
return helper.get(union_name)
|
||||
|
||||
def new_cond_fn(*args):
|
||||
"""cond is a zero-args function, which is not
|
||||
compatible with `while_loop`.
|
||||
"""
|
||||
return cond()
|
||||
|
||||
# UndefinedVar will become data layer not check variable with value=NO_VALUE_MAGIC.
|
||||
loop_vars = [
|
||||
to_static_variable(var) if not isinstance(var, UndefinedVar) else var
|
||||
for var in helper.get(union_name)
|
||||
]
|
||||
helper.set(union_name, loop_vars) # change the non-local var to variable
|
||||
# variable maybe modified to inner var. change it into
|
||||
from paddle.static.nn import while_loop
|
||||
|
||||
loop_vars = while_loop(new_cond_fn, new_body_fn, loop_vars)
|
||||
helper.set(union_name, loop_vars)
|
||||
return loop_vars
|
||||
|
||||
|
||||
def _run_py_while(cond, body, getter, setter):
|
||||
while True:
|
||||
pred = cond()
|
||||
if isinstance(pred, (Variable, Value)):
|
||||
raise Dygraph2StaticException(
|
||||
"python while pred change from bool to variable."
|
||||
)
|
||||
if not pred:
|
||||
break
|
||||
body()
|
||||
|
||||
|
||||
def convert_logical_and(x_func, y_func):
|
||||
"""
|
||||
A function representation of a Python ``and`` statement.
|
||||
|
||||
Args:
|
||||
x_func(callable): x_func() is the left hand operand of ``and`` operator. x_func() is bool or Tensor.
|
||||
y_func(callable): y_func() is the right hand operand of ``and`` operator. y_func() is bool or Tensor.
|
||||
|
||||
Returns:
|
||||
A python bool variable or a bool Tensor.
|
||||
|
||||
NOTE(liym27):
|
||||
1) The operands are executed sequentially according to the running logic of Python. So here the arguments
|
||||
should be callable.
|
||||
2) If the left hand operand is False, the right hand operand should be executed.
|
||||
|
||||
For example:
|
||||
a = x > 1 and y < 1
|
||||
Transformed code:
|
||||
a = paddle.jit.dy2static.convert_logical_and(lambda:x>1, lambda:y<1)
|
||||
|
||||
In `convert_logical_and(lambda:x>1, lambda:y<1)`, `lambda:y<1` must be run after `lambda:x>1`. And
|
||||
if `x>1` is False, `y<1` should NOT be run.
|
||||
"""
|
||||
x_value = x_func()
|
||||
if not isinstance(x_value, (Variable, Value)):
|
||||
return _run_py_logical_and(lambda: x_value, y_func)
|
||||
|
||||
y_value = y_func()
|
||||
if not isinstance(y_value, (Variable, Value)):
|
||||
return _run_py_logical_and(lambda: y_value, lambda: x_value)
|
||||
|
||||
return _run_paddle_logical_and(x_value, y_value)
|
||||
|
||||
|
||||
def _run_paddle_logical_and(x, y):
|
||||
x = cast_bool_if_necessary(x)
|
||||
y = cast_bool_if_necessary(y)
|
||||
return paddle.logical_and(x, y)
|
||||
|
||||
|
||||
def _run_py_logical_and(x_func, y_func):
|
||||
x_value = x_func()
|
||||
assert not isinstance(x_value, (Variable, Value))
|
||||
|
||||
# NOTE(liym27):
|
||||
# 1. Returns y_func() if x_value is False;
|
||||
# 2. If x_value is False, y_func() should not be run.
|
||||
return x_value and y_func()
|
||||
|
||||
|
||||
def convert_logical_or(x_func, y_func):
|
||||
"""
|
||||
A function representation of a Python ``or`` statement.
|
||||
|
||||
Args:
|
||||
x_func(callable): x_func() is the left hand operand of ``or`` operator. x_func() is bool or Tensor.
|
||||
y_func(callable): y_func() is the right hand operand of ``or`` operator. y_func() is bool or Tensor.
|
||||
|
||||
Returns:
|
||||
A python bool variable or a bool Tensor.
|
||||
|
||||
NOTE(liym27):
|
||||
1) The operands are executed sequentially according to the running logic of Python. So here the arguments
|
||||
should be callable.
|
||||
2) If the left hand operand is True, the right hand operand should be executed.
|
||||
|
||||
For example:
|
||||
a = x > 1 or y < 1
|
||||
Transformed code:
|
||||
a = paddle.jit.dy2static.convert_logical_or(lambda:x>1, lambda:y<1)
|
||||
|
||||
In `convert_logical_or(lambda:x>1, lambda:y<1)`, `lambda:y<1` must be run after `lambda:x>1`. And
|
||||
if `x>1` is True, `y<1` should NOT be run.
|
||||
"""
|
||||
x_value = x_func()
|
||||
if not isinstance(x_value, (Variable, Value)):
|
||||
return _run_py_logical_or(lambda: x_value, y_func)
|
||||
|
||||
y_value = y_func()
|
||||
if not isinstance(y_value, (Variable, Value)):
|
||||
return _run_py_logical_or(lambda: y_value, lambda: x_value)
|
||||
|
||||
return _run_paddle_logical_or(x_value, y_value)
|
||||
|
||||
|
||||
def _run_paddle_logical_or(x, y):
|
||||
x = cast_bool_if_necessary(x)
|
||||
y = cast_bool_if_necessary(y)
|
||||
return paddle.logical_or(x, y)
|
||||
|
||||
|
||||
def _run_py_logical_or(x_func, y_func):
|
||||
x_value = x_func()
|
||||
assert not isinstance(x_value, (Variable, Value))
|
||||
|
||||
# NOTE(liym27):
|
||||
# 1. Returns y_func() if x_value is False;
|
||||
# 2. If x_value is True, y_func() should not be run.
|
||||
return x_value or y_func()
|
||||
|
||||
|
||||
def convert_logical_not(x):
|
||||
"""
|
||||
A function representation of a Python ``not`` statement.
|
||||
|
||||
Args:
|
||||
x(bool|Tensor): Operand of ``not`` operator.
|
||||
|
||||
Returns:
|
||||
A python bool variable or a bool Tensor.
|
||||
"""
|
||||
|
||||
if isinstance(x, (Variable, Value)):
|
||||
return _run_paddle_logical_not(x)
|
||||
else:
|
||||
return _run_py_logical_not(x)
|
||||
|
||||
|
||||
def _run_paddle_logical_not(x):
|
||||
x = cast_bool_if_necessary(x)
|
||||
return paddle.logical_not(x)
|
||||
|
||||
|
||||
def _run_py_logical_not(x):
|
||||
return not x
|
||||
|
||||
|
||||
def convert_ifelse(
|
||||
pred,
|
||||
true_fn,
|
||||
false_fn,
|
||||
get_args,
|
||||
set_args,
|
||||
return_name_ids,
|
||||
push_pop_names=None,
|
||||
):
|
||||
"""
|
||||
A function representation of a Python ``if/else`` statement.
|
||||
|
||||
Args:
|
||||
pred(bool|Tensor): A boolean Tensor which determines whether to return the result of ``true_fn`` or ``false_fn`` .
|
||||
true_fn(callable): A callable to be performed if ``pred`` is true.
|
||||
false_fn(callable): A callable to be performed if ``pred`` is false.
|
||||
get_args(callable): Get all arguments that needed in true_fn and false_fn.
|
||||
set_args(callable): Update arguments that modified in trure_fn and false_fn.
|
||||
return_name_ids(list[string], optional): the returned names.
|
||||
push_pop_names(list[string], optional): the names on which called .append() or .pop().
|
||||
|
||||
Returns:
|
||||
``true_fn()`` if the predicate ``pred`` is true else ``false_fn()`` .
|
||||
|
||||
"""
|
||||
if isinstance(pred, (Variable, Value)):
|
||||
out = _run_paddle_cond(
|
||||
pred,
|
||||
true_fn,
|
||||
false_fn,
|
||||
get_args,
|
||||
set_args,
|
||||
return_name_ids,
|
||||
push_pop_names,
|
||||
)
|
||||
else:
|
||||
out = _run_py_ifelse(
|
||||
pred, true_fn, false_fn, get_args, set_args, return_name_ids
|
||||
)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def _run_paddle_cond(
|
||||
pred, true_fn, false_fn, get_args, set_args, return_name_ids, push_pop_names
|
||||
):
|
||||
"""
|
||||
Paddle cond API will evaluate both true_fn and false_fn codes.
|
||||
"""
|
||||
helper = GetterSetterHelper(
|
||||
get_args, set_args, return_name_ids, push_pop_names
|
||||
)
|
||||
_convert_tensor_array_if_necessary(helper, push_pop_names)
|
||||
pred = cast_bool_if_necessary(pred)
|
||||
init_args = helper.get(return_name_ids)
|
||||
from paddle.jit.dy2static.parameter_recorder import _global_inplace_map
|
||||
from paddle.jit.dy2static.program_translator import ProgramTranslator
|
||||
|
||||
if use_pir_api():
|
||||
inplace_map = _global_inplace_map
|
||||
else:
|
||||
inplace_map = ProgramTranslator.get_instance()._inplace_map
|
||||
union_name = None
|
||||
# TODO(@xiongkun) lambda can have push_pop_names, which will cause error.
|
||||
if return_name_ids is None and push_pop_names is None:
|
||||
union_name = None
|
||||
else:
|
||||
union_name = (
|
||||
OrderedSet(return_name_ids) if return_name_ids else OrderedSet()
|
||||
) | (OrderedSet(push_pop_names) if push_pop_names else OrderedSet())
|
||||
union_name = list(union_name)
|
||||
|
||||
def new_true_fn():
|
||||
nonlocal union_name
|
||||
# init args may contain mutable python container like [var, 2], we copy then like in while_loop
|
||||
inplace_map_checkpoint = inplace_map.save_checkpoint()
|
||||
helper.set(
|
||||
return_name_ids,
|
||||
paddle.utils.copy_mutable_vars(init_args),
|
||||
)
|
||||
ret = true_fn()
|
||||
# IfExpr will return a non-None return value, so we just return ret.
|
||||
# We assume normal return has no return value.
|
||||
if ret is None:
|
||||
ret = helper.get(union_name)
|
||||
inplace_map.restore_checkpoint(inplace_map_checkpoint)
|
||||
return ret
|
||||
|
||||
def new_false_fn():
|
||||
nonlocal union_name
|
||||
# init args may contain mutable python container like [var, 2], we copy then like in while_loop
|
||||
inplace_map_checkpoint = inplace_map.save_checkpoint()
|
||||
helper.set(
|
||||
return_name_ids,
|
||||
paddle.utils.copy_mutable_vars(init_args),
|
||||
)
|
||||
ret = false_fn()
|
||||
if ret is None:
|
||||
ret = helper.get(union_name)
|
||||
inplace_map.restore_checkpoint(inplace_map_checkpoint)
|
||||
return ret
|
||||
|
||||
try:
|
||||
cond_outs = paddle.static.nn.cond(
|
||||
pred, new_true_fn, new_false_fn, None, union_name
|
||||
)
|
||||
except Exception as e:
|
||||
if re.search(
|
||||
"Unsupported return type of true_fn and false_fn in cond", str(e)
|
||||
):
|
||||
raise Dygraph2StaticException(
|
||||
f"Your if/else have different return type. TODO: add link to modify. {e}"
|
||||
)
|
||||
if re.search("Incompatible return values of", str(e)):
|
||||
raise Dygraph2StaticException(
|
||||
f"Your if/else have different number of return value. TODO: add link to modify. {e}"
|
||||
)
|
||||
raise e
|
||||
get_args = lambda: helper.get(union_name)
|
||||
set_args = lambda vs: helper.set(union_name, vs)
|
||||
return _recover_args_state(cond_outs, get_args, set_args, union_name)
|
||||
|
||||
|
||||
def _run_py_ifelse(
|
||||
pred, true_fn, false_fn, get_args, set_args, return_name_ids
|
||||
):
|
||||
"""
|
||||
Evaluate python original branch function if-else.
|
||||
"""
|
||||
py_outs = true_fn() if pred else false_fn()
|
||||
return py_outs
|
||||
|
||||
|
||||
def _remove_no_value_return_var(out):
|
||||
if isinstance(out, tuple) and len(out) > 0:
|
||||
processed_out = out
|
||||
align_ret = out[0]
|
||||
if isinstance(align_ret, tuple):
|
||||
for index, item in enumerate(align_ret):
|
||||
if isinstance(item, (Variable, Value)) and (
|
||||
RETURN_NO_VALUE_VAR_NAME in item.name
|
||||
):
|
||||
# return None
|
||||
if index == 0:
|
||||
processed_out = (None, *out[1:])
|
||||
elif index == 1:
|
||||
processed_out = align_ret[:1] + out[1:]
|
||||
else:
|
||||
processed_out = (align_ret[:index], *out[1:])
|
||||
break
|
||||
|
||||
for index, item in enumerate(processed_out):
|
||||
if isinstance(item, (Variable, Value)) and (
|
||||
RETURN_NO_VALUE_VAR_NAME in item.name
|
||||
):
|
||||
processed_out = processed_out[:index]
|
||||
|
||||
if not processed_out:
|
||||
return None
|
||||
elif len(processed_out) == 1:
|
||||
return processed_out[0]
|
||||
else:
|
||||
return processed_out
|
||||
|
||||
else:
|
||||
return out
|
||||
|
||||
|
||||
def _check_no_undefined_var(outs, names, branch_name):
|
||||
if names is None:
|
||||
return
|
||||
if not isinstance(outs, (list, tuple)):
|
||||
outs = [outs]
|
||||
for var, name in zip(list(outs), names):
|
||||
if isinstance(var, UndefinedVar):
|
||||
raise ValueError(
|
||||
f"Required '{name}' must be initialized both in if-else branch, but found it not initialized in '{branch_name}'."
|
||||
)
|
||||
|
||||
|
||||
def _recover_args_state(outs, get_args, set_args, return_name_ids):
|
||||
"""
|
||||
Currently we support variant length of early return statement by padding
|
||||
_no_return_value.
|
||||
|
||||
# TODO(dev): We shall consider to evaluate whether should support this for Python if-else?
|
||||
"""
|
||||
# IfExpr's return_name_ids maybe None
|
||||
if return_name_ids is None:
|
||||
return outs
|
||||
|
||||
init_args = get_args()
|
||||
# recover args state
|
||||
num_outs = len(return_name_ids)
|
||||
num_args = len(init_args)
|
||||
assert num_outs <= num_args
|
||||
|
||||
if num_args == 1:
|
||||
final_outs = (
|
||||
(outs,) if not isinstance(outs, (list, tuple)) else tuple(outs)
|
||||
)
|
||||
else:
|
||||
outs = (outs,) if num_outs == 1 else tuple(outs)
|
||||
final_outs = outs + init_args[num_outs:]
|
||||
|
||||
set_args(final_outs)
|
||||
return final_outs
|
||||
|
||||
|
||||
def convert_len(var):
|
||||
"""
|
||||
Returns variable(length) from shape ops based on var.type
|
||||
|
||||
Note: In addition to some ast transformations, some block-related
|
||||
operations are added in `len` transformation, such as appending
|
||||
`shape_op` in var.block.
|
||||
"""
|
||||
if isinstance(var, Variable):
|
||||
assert var.ndim > 0, "len() of a 0-D tensor is wrong"
|
||||
if var.type in [
|
||||
core.VarDesc.VarType.DENSE_TENSOR,
|
||||
core.VarDesc.VarType.SELECTED_ROWS,
|
||||
]:
|
||||
# Note: Length of var may be known ahead of time in dygraph,
|
||||
# but it probably represents batch size which can be variant.
|
||||
# so we return a variable dynamically inferred from var.shape.
|
||||
if (
|
||||
var.shape[0] > 0
|
||||
and var.type == core.VarDesc.VarType.DENSE_TENSOR
|
||||
):
|
||||
return var.shape[0]
|
||||
return paddle.shape(var)[0]
|
||||
elif var.type == core.VarDesc.VarType.DENSE_TENSOR_ARRAY:
|
||||
return paddle.tensor.array_length(var)
|
||||
else:
|
||||
raise TypeError(
|
||||
f'len(var) only supports DenseTensor/DenseTensorArray/SelectedRows, but received {type(var)}.'
|
||||
)
|
||||
elif isinstance(var, Value):
|
||||
if var.is_dense_tensor_type() or var.is_selected_row_type():
|
||||
assert var.ndim > 0, "len() of a 0-D tensor is wrong"
|
||||
# Note: Length of var may be known ahead of time in dygraph,
|
||||
# but it probably represents batch size which can be variant.
|
||||
# so we return a variable dynamically inferred from var.shape.
|
||||
if var.shape[0] > 0 and var.is_dense_tensor_type():
|
||||
return var.shape[0]
|
||||
return paddle.shape(var)[0]
|
||||
elif var.is_dense_tensor_array_type():
|
||||
return paddle.tensor.array_length(var)
|
||||
else:
|
||||
raise TypeError(
|
||||
'len(var) only supports DenseTensor/DenseTensorArray/SelectedRows, '
|
||||
+ f'but received {type(var)}.'
|
||||
)
|
||||
else:
|
||||
if isinstance(var, VariableTuple):
|
||||
return var.__len__()
|
||||
return len(var)
|
||||
|
||||
|
||||
def convert_zip(*args):
|
||||
for i, arg in enumerate(args):
|
||||
if isinstance(arg, (Variable, Value)) and arg.shape[0] == -1:
|
||||
raise RuntimeError(
|
||||
"Not support zip(tensor, ...) when tensor.shape[0] == -1, "
|
||||
f"but found args[{i}].shape[0] == -1 in 'zip'"
|
||||
)
|
||||
return zip(*args)
|
||||
|
||||
|
||||
def convert_super(super_fn):
|
||||
if super_fn is super:
|
||||
return super_fn
|
||||
return lambda cls, instance: super_fn()
|
||||
|
||||
|
||||
# TODO(xiongkun): delete when list<variable> is ready.
|
||||
class VariableTuple:
|
||||
"""
|
||||
this class will cause enumerate can't be wrapped by other iterator change function.
|
||||
this will be fixed when list<Variable> is produced.
|
||||
VariableTuple can only deal with variables which is fixed.
|
||||
"""
|
||||
|
||||
def __init__(self, var, start=0):
|
||||
self.var = var
|
||||
self.len = convert_len(var)
|
||||
if isinstance(self.len, (Variable, Value)):
|
||||
self.rag = paddle.arange(start, start + self.len, 1, "int64")
|
||||
else:
|
||||
self.rag = range(start, start + self.len)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
return self.rag[idx], self.var[idx]
|
||||
|
||||
def __len__(self):
|
||||
return self.len
|
||||
|
||||
|
||||
def convert_enumerate(*args):
|
||||
has_variable = any(isinstance(x, (Variable, Value)) for x in args)
|
||||
if has_variable:
|
||||
return VariableTuple(*args)
|
||||
return enumerate(*args)
|
||||
|
||||
|
||||
def convert_range(*args):
|
||||
has_variable = any(isinstance(x, (Variable, Value)) for x in args)
|
||||
# NOTE(SigureMo): Add an `Assign` OP after the Tensor input to mark it as a variable, which can
|
||||
# avoid confusing it with the scalar case in `arange` API.
|
||||
# For example:
|
||||
# ```python
|
||||
# l = []
|
||||
# for i in range(n):
|
||||
# l.append(i)
|
||||
# ```
|
||||
# - If `n` is a scalar (e.g., `n=10`), we expect to create an `ArangeOp` with a fixed output shape [10].
|
||||
# - If `n` is a Tensor (e.g., `n=full([], 10, "int32")`), we expect to create an `ArangeOp` with a dynamic
|
||||
# output shape [-1]. To ensure the python level and graph level all recognize this is data-dependent control
|
||||
# flow.
|
||||
# However, we can't distinguish the scalar case and the Tensor case when creating the `ArangeOp`. Because
|
||||
# the scalar case also be convert as a `Full` OP output.
|
||||
# So we add an `Assign` OP after the Tensor input to **mark** it as a variable, which can avoid confusing
|
||||
# it with the scalar case.
|
||||
is_full_op_output = lambda x: (
|
||||
isinstance(x, Value)
|
||||
and x.get_defining_op()
|
||||
and x.get_defining_op().name() == "pd_op.full"
|
||||
)
|
||||
args = [
|
||||
paddle.assign(arg) if is_full_op_output(arg) else arg for arg in args
|
||||
]
|
||||
if has_variable:
|
||||
if len(args) == 1:
|
||||
return paddle.arange(0, args[0], 1, "int64")
|
||||
if len(args) == 2:
|
||||
return paddle.arange(args[0], args[1], 1, "int64")
|
||||
if len(args) == 3:
|
||||
return paddle.arange(args[0], args[1], args[2], "int64")
|
||||
return range(*args)
|
||||
|
||||
|
||||
def convert_shape(x):
|
||||
"""
|
||||
A function representation of the shape of variable.
|
||||
"""
|
||||
|
||||
def has_negative(list_shape):
|
||||
return any(x < 0 for x in list_shape)
|
||||
|
||||
# When `x` is Variable:
|
||||
# (1) if x.shape contains -1, such as [2, -1, 64], returns [2, var, 64],
|
||||
# where var = paddle.shape(x)[1]
|
||||
|
||||
# (2) if x.shape does not contains -1, return list(x.shape) directly
|
||||
|
||||
if isinstance(x, (Variable, Value)):
|
||||
values = list(x.shape)
|
||||
if has_negative(values):
|
||||
shape_tensor = paddle.shape(x)
|
||||
for i, v in enumerate(values):
|
||||
if v is None or v < 0:
|
||||
values[i] = shape_tensor[i]
|
||||
return values
|
||||
else:
|
||||
return x.shape
|
||||
|
||||
|
||||
def cast_bool_if_necessary(var):
|
||||
assert isinstance(var, (Variable, Value))
|
||||
if convert_dtype(var.dtype) not in ['bool']:
|
||||
var = paddle.cast(var, dtype="bool")
|
||||
return var
|
||||
|
||||
|
||||
def convert_var_dtype(var, dtype):
|
||||
if isinstance(var, (Variable, Value)):
|
||||
src_dtype = convert_dtype(var.dtype)
|
||||
assert src_dtype in [
|
||||
'bool',
|
||||
'float16',
|
||||
'float32',
|
||||
'float64',
|
||||
'int32',
|
||||
'int64',
|
||||
'uint8',
|
||||
], (
|
||||
f"The dtype of var {var.name} is {src_dtype}, which is not supported in the cast op."
|
||||
)
|
||||
assert dtype in [
|
||||
'bool',
|
||||
'int',
|
||||
'float',
|
||||
'complex',
|
||||
], (
|
||||
f"The casted target dtype is {dtype}, which is not supported in type casting."
|
||||
)
|
||||
cast_map = {
|
||||
'bool': 'bool',
|
||||
'int': 'int32',
|
||||
'float': 'float32',
|
||||
'complex': 'complex64',
|
||||
}
|
||||
return paddle.cast(var, dtype=cast_map[dtype])
|
||||
else:
|
||||
assert dtype in [
|
||||
'bool',
|
||||
'int',
|
||||
'float',
|
||||
'complex',
|
||||
], (
|
||||
f"The casted target dtype is {dtype}, which is not supported in type casting."
|
||||
)
|
||||
return eval(dtype)(var)
|
||||
|
||||
|
||||
def convert_assert(cond, message=""):
|
||||
"""
|
||||
A function representation of a Python ``assert`` statement.
|
||||
"""
|
||||
if isinstance(cond, (Variable, Value)):
|
||||
cond = paddle.cast(cond, "bool")
|
||||
# NOTE: message is not used because Paddle Assert has no corresponding parameter to use.
|
||||
from paddle.static.nn.control_flow import Assert
|
||||
|
||||
return Assert(cond)
|
||||
else:
|
||||
assert cond, message
|
||||
|
||||
|
||||
def convert_print(*objects, sep=' ', end='\n', file=None, flush=False):
|
||||
"""
|
||||
A function representing Python ``print`` function. It will print all arguments
|
||||
at compile time and only print the Tensor values at runtime.
|
||||
"""
|
||||
for obj in objects:
|
||||
if isinstance(obj, (Variable, Value)):
|
||||
paddle.static.Print(obj)
|
||||
print(*objects, sep=sep, end=end, file=file, flush=flush)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def convert_auto_cast(
|
||||
enable=True,
|
||||
custom_white_list=None,
|
||||
custom_black_list=None,
|
||||
level='O1',
|
||||
dtype='float16',
|
||||
use_promote=True,
|
||||
):
|
||||
from .program_translator import ProgramTranslator
|
||||
|
||||
warnings.warn(
|
||||
"paddle.amp.auto_cast is an experimental features in auto parallel."
|
||||
+ "This will take no effect in normal dy2static."
|
||||
)
|
||||
|
||||
amp_records = ProgramTranslator.get_instance()._amp_records
|
||||
main_program = paddle.static.default_main_program()
|
||||
current_block_idx = main_program.current_block_idx
|
||||
current_block = main_program.current_block()
|
||||
start_op_idx = len(current_block.ops)
|
||||
amp_options = AmpOptions(
|
||||
enable, custom_white_list, custom_black_list, level, dtype, use_promote
|
||||
)
|
||||
yield
|
||||
end_op_idx = len(current_block.ops)
|
||||
if current_block_idx not in amp_records:
|
||||
amp_records[current_block_idx] = []
|
||||
amp_records[current_block_idx].append(
|
||||
(amp_options, start_op_idx, end_op_idx)
|
||||
)
|
||||
|
||||
|
||||
def create_bool_as_type(x, value=True):
|
||||
'''
|
||||
Create a bool variable, which type is the same as x.
|
||||
'''
|
||||
if isinstance(x, (Variable, Value)):
|
||||
return paddle.full(shape=[], fill_value=value, dtype="bool")
|
||||
else:
|
||||
return value
|
||||
Reference in New Issue
Block a user