def elementwise_binary(alpha, A, desc_A, mode_A, gamma, C, desc_C, mode_C, out=None, op_AC=cutensor.OP_ADD, compute_dtype=None): """Element-wise tensor operation for two input tensors This function performs a element-wise tensor operation of the form: D_{Pi^C(i_0,i_1,...,i_n)} = op_AC(alpha * uop_A(A_{Pi^A(i_0,i_1,...,i_n)}), gamma * uop_C(C_{Pi^C(i_0,i_1,...,i_n)})) See elementwise_trinary() for details. Examples: See examples/cutensor/elementwise_binary.py """ if A.dtype != C.dtype: raise ValueError('dtype mismatch: {} != {}'.format(A.dtype, C.dtype)) if not (A.flags.c_contiguous and C.flags.c_contiguous): raise ValueError('The inputs should be contiguous arrays.') if out is None: out = cupy.ndarray(C.shape, dtype=C.dtype) elif C.dtype != out.dtype: raise ValueError('dtype mismatch: {} != {}'.format(C.dtype, out.dtype)) elif C.shape != out.shape: raise ValueError('shape mismatch: {} != {}'.format(C.shape, out.shape)) elif not out.flags.c_contiguous: raise ValueError('`out` should be a contiguous array.') if A.ndim != len(mode_A): raise ValueError('ndim mismatch: {} != {}'.format(A.ndim, len(mode_A))) if C.ndim != len(mode_C): raise ValueError('ndim mismatch: {} != {}'.format(C.ndim, len(mode_C))) mode_A = _convert_mode(mode_A) mode_C = _convert_mode(mode_C) if compute_dtype is None: compute_dtype = A.dtype alpha = numpy.array(alpha, compute_dtype) gamma = numpy.array(gamma, compute_dtype) handle = get_handle() cuda_dtype = get_cuda_dtype(compute_dtype) cutensor.elementwiseBinary(handle, alpha.ctypes.data, A.data.ptr, desc_A, mode_A.ctypes.data, gamma.ctypes.data, C.data.ptr, desc_C, mode_C.ctypes.data, out.data.ptr, desc_C, mode_C.ctypes.data, op_AC, cuda_dtype) return out
def elementwise_binary(alpha, A, desc_A, mode_A, gamma, C, desc_C, mode_C, out=None, op_AC=cutensor.OP_ADD, compute_dtype=None): """Element-wise tensor operation for two input tensors This function performs a element-wise tensor operation of the form: D_{Pi^C(i_0,i_1,...,i_n)} = op_AC(alpha * uop_A(A_{Pi^A(i_0,i_1,...,i_n)}), gamma * uop_C(C_{Pi^C(i_0,i_1,...,i_n)})) See elementwise_trinary() for details. Examples: See examples/cutensor/elementwise_binary.py """ assert A.dtype == C.dtype assert A.ndim == len(mode_A) assert C.ndim == len(mode_C) mode_A = _convert_mode(mode_A) mode_C = _convert_mode(mode_C) if out is None: out = cupy.ndarray(C.shape, dtype=C.dtype) else: assert C.dtype == out.dtype assert C.ndim == out.ndim for i in range(C.ndim): assert C.shape[i] == out.shape[i] if compute_dtype is None: compute_dtype = A.dtype alpha = numpy.array(alpha, compute_dtype) gamma = numpy.array(gamma, compute_dtype) handle = get_handle() cuda_dtype = get_cuda_dtype(compute_dtype) cutensor.elementwiseBinary(handle, alpha.ctypes.data, A.data.ptr, desc_A, mode_A.ctypes.data, gamma.ctypes.data, C.data.ptr, desc_C, mode_C.ctypes.data, out.data.ptr, desc_C, mode_C.ctypes.data, op_AC, cuda_dtype) return out