# -*- coding: utf-8 -*-
#
#
# Unless required by applicable law or agreed to in writing,
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# pylint: disable=redefined-builtin
from typing import Iterable, Union

from ..core._imperative_rt.core2 import pop_scope, push_scope
from ..functional import clip, concat, minimum, norm
from ..tensor import Tensor

tensors: Union[Tensor, Iterable[Tensor]], max_norm: float, ord: float = 2.0,
):
r"""Clips gradient norm of an iterable of parameters.
The norm is computed over all gradients together, as if they were
concatenated into a single vector. Gradients are modified in-place.

Args:
tensors: an iterable of Tensors or a single Tensor.
max_norm: max norm of the gradients.
ord: type of the used p-norm. Can be 'inf' for infinity norm.

Returns:
total norm of the parameters (viewed as a single vector).
"""
if isinstance(tensors, Tensor):
tensors = [tensors]
tensors = [t for t in tensors if t.grad is not None]
if len(tensors) == 0:
return Tensor(0.0)
norm_ = [norm(t.grad.flatten(), ord=ord) for t in tensors]
if len(norm_) > 1:
norm_ = norm(concat(norm_), ord=ord)
else:
norm_ = norm_
scale = max_norm / (norm_ + 1e-6)
scale = minimum(scale, 1)
for tensor in tensors:
return norm_

tensors: Union[Tensor, Iterable[Tensor]], lower: float, upper: float
):
r"""Clips gradient of an iterable of parameters to a specified lower and

The gradients are clipped in the range:

.. math:: \left[\text{lower}, \text{upper}\right]

Args:
tensors: an iterable of Tensors or a single Tensor.
lower: minimum allowed value of the gradients.
upper: maximum allowed value of the gradients.
"""