# Copyright (c) 2017 Sony Corporation. 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.

"""
This is still left unlike other *2 operation (sub2, mul2, ...) because
it has cudnn implementation.
"""
import pytest
import numpy as np
import nnabla as nn
import nnabla.functions as F
from nbla_test_utils import list_context

ctxs = list_context('Assign')


@pytest.mark.parametrize("ctx, func_name", ctxs)
@pytest.mark.parametrize("seed", [314])
def test_assign_forward_backward(seed, ctx, func_name):
    from nbla_test_utils import function_tester
    rng = np.random.RandomState(seed)
    inputs = [rng.randn(2, 3, 4).astype(np.float32) * 2 for _ in range(2)]
    grads = np.zeros((48,))
    function_tester(rng, F.assign, lambda dst, src: src, inputs, ref_grad=lambda *args: grads,
                    ctx=ctx, func_name=func_name, atol_f=1e-3, atol_b=1e-2)

@@ -1017,360 +1017,533 @@
     
-from nbla_test_utils import function_tester%0A    rng = np.random.RandomState(seed)%0A    inputs = %5Brng.randn(2, 3, 4).astype(np.float32) * 2 for _ in range(2)%5D%0A    grads = np.zeros((48,))%0A    function_tester(rng, F.assign, lambda 
+rng = np.random.RandomState(seed)%0A    dst = nn.Variable((2, 3, 4), need_grad=True)%0A    src = nn.Variable((2, 3, 4), need_grad=True)%0A%0A    assign = F.assign(dst, src)%0A%0A    src.d = np.random.random((2, 3, 4))%0A    assign.forward()%0A%0A    assert np.allclose(
 dst
+.d
 , src
-: src, inputs, ref_grad=lambda *args: grads,%0A                    ctx=ctx, func_name=func_name, atol_f=1e-3, atol_b=1e-2
+.d)%0A%0A    dummy = assign + np.random.random()%0A%0A    dst.grad.zero()%0A    src.grad.zero()%0A    dummy.forward()%0A    dummy.backward()%0A%0A    # assign should not propagate gradients%0A    assert np.all(dst.g == np.zeros((2, 3, 4)))%0A    assert np.all(src.g == np.zeros((2, 3, 4))
 )%0A

Update assign test