from chainer import cuda
from chainer import function


class Copy(function.Function):

    """Copy an input GPUArray onto another device."""

    def __init__(self, out_device):
        self.out_device = out_device

    def forward_cpu(self, x):
        return x[0].copy(),

    def forward_gpu(self, x):
        return cuda.copy(x[0], out_device=self.out_device),

    def backward_cpu(self, x, gy):
        return gy[0].copy(),

    def backward_gpu(self, x, gy):
        return cuda.copy(gy[0], out_device=cuda.get_device(x[0])),


def copy(x, dst):
    """Copies the input variable onto the specified device.

    This function copies the array of input variable onto the device specified
    by ``dst`` if the original array is on GPU, and otherwise just copies the
    array within host memory.

    Args:
        x (~chainer.Variable): Variable to be copied.
        dst: Target device specifier.

    Returns:
        ~chainer.Variable: Output variable.

    """
    return Copy(dst)(x)

import numpy

from chainer import cuda
from chainer import function
from chainer.utils import type_check


class Copy(function.Function):

    """Copy an input GPUArray onto another device."""

    def __init__(self, out_device):
        self.out_device = out_device

    def check_type_forward(self, in_types):
        type_check.expect(
            in_types.size() == 1,
            in_types[0].dtype == numpy.float32
        )

    def check_type_backward(self, in_types, out_types):
        type_check.expect(
            out_types.size() == 1,
            in_types[0].dtype == out_types[0].dtype,
            in_types[0].ndim == out_types[0].ndim,
            in_types[0].shape == out_types[0].shape
        )

    def forward_cpu(self, x):
        return x[0].copy(),

    def forward_gpu(self, x):
        return cuda.copy(x[0], out_device=self.out_device),

    def backward_cpu(self, x, gy):
        return gy[0].copy(),

    def backward_gpu(self, x, gy):
        return cuda.copy(gy[0], out_device=cuda.get_device(x[0])),


def copy(x, dst):
    """Copies the input variable onto the specified device.

    This function copies the array of input variable onto the device specified
    by ``dst`` if the original array is on GPU, and otherwise just copies the
    array within host memory.

    Args:
        x (~chainer.Variable): Variable to be copied.
        dst: Target device specifier.

    Returns:
        ~chainer.Variable: Output variable.

    """
    return Copy(dst)(x)

Add unittest(cpu-only) and typecheck for Copy
Add unittest(cpu-only) and typecheck for Copy

+ import numpy
+ 
  from chainer import cuda
  from chainer import function
+ from chainer.utils import type_check
  
  
  class Copy(function.Function):
  
      """Copy an input GPUArray onto another device."""
  
      def __init__(self, out_device):
          self.out_device = out_device
+ 
+     def check_type_forward(self, in_types):
+         type_check.expect(
+             in_types.size() == 1,
+             in_types[0].dtype == numpy.float32
+         )
+ 
+     def check_type_backward(self, in_types, out_types):
+         type_check.expect(
+             out_types.size() == 1,
+             in_types[0].dtype == out_types[0].dtype,
+             in_types[0].ndim == out_types[0].ndim,
+             in_types[0].shape == out_types[0].shape
+         )
  
      def forward_cpu(self, x):
          return x[0].copy(),
  
      def forward_gpu(self, x):
          return cuda.copy(x[0], out_device=self.out_device),
  
      def backward_cpu(self, x, gy):
          return gy[0].copy(),
  
      def backward_gpu(self, x, gy):
          return cuda.copy(gy[0], out_device=cuda.get_device(x[0])),
  
  
  def copy(x, dst):
      """Copies the input variable onto the specified device.
  
      This function copies the array of input variable onto the device specified
      by ``dst`` if the original array is on GPU, and otherwise just copies the
      array within host memory.
  
      Args:
          x (~chainer.Variable): Variable to be copied.
          dst: Target device specifier.
  
      Returns:
          ~chainer.Variable: Output variable.
  
      """
      return Copy(dst)(x)
  
Add unittest(cpu-only) and typecheck for Copy
## Code Before:
from chainer import cuda
from chainer import function


class Copy(function.Function):

    """Copy an input GPUArray onto another device."""

    def __init__(self, out_device):
        self.out_device = out_device

    def forward_cpu(self, x):
        return x[0].copy(),

    def forward_gpu(self, x):
        return cuda.copy(x[0], out_device=self.out_device),

    def backward_cpu(self, x, gy):
        return gy[0].copy(),

    def backward_gpu(self, x, gy):
        return cuda.copy(gy[0], out_device=cuda.get_device(x[0])),


def copy(x, dst):
    """Copies the input variable onto the specified device.

    This function copies the array of input variable onto the device specified
    by ``dst`` if the original array is on GPU, and otherwise just copies the
    array within host memory.

    Args:
        x (~chainer.Variable): Variable to be copied.
        dst: Target device specifier.

    Returns:
        ~chainer.Variable: Output variable.

    """
    return Copy(dst)(x)

## Instruction:
Add unittest(cpu-only) and typecheck for Copy
## Code After:
import numpy

from chainer import cuda
from chainer import function
from chainer.utils import type_check


class Copy(function.Function):

    """Copy an input GPUArray onto another device."""

    def __init__(self, out_device):
        self.out_device = out_device

    def check_type_forward(self, in_types):
        type_check.expect(
            in_types.size() == 1,
            in_types[0].dtype == numpy.float32
        )

    def check_type_backward(self, in_types, out_types):
        type_check.expect(
            out_types.size() == 1,
            in_types[0].dtype == out_types[0].dtype,
            in_types[0].ndim == out_types[0].ndim,
            in_types[0].shape == out_types[0].shape
        )

    def forward_cpu(self, x):
        return x[0].copy(),

    def forward_gpu(self, x):
        return cuda.copy(x[0], out_device=self.out_device),

    def backward_cpu(self, x, gy):
        return gy[0].copy(),

    def backward_gpu(self, x, gy):
        return cuda.copy(gy[0], out_device=cuda.get_device(x[0])),


def copy(x, dst):
    """Copies the input variable onto the specified device.

    This function copies the array of input variable onto the device specified
    by ``dst`` if the original array is on GPU, and otherwise just copies the
    array within host memory.

    Args:
        x (~chainer.Variable): Variable to be copied.
        dst: Target device specifier.

    Returns:
        ~chainer.Variable: Output variable.

    """
    return Copy(dst)(x)
