yield vect def randn(self): _randn = (np.array(np.random.randn(*self.shape)).astype(self.dtype) + 1.0j * np.array(np.random.randn(*self.shape)).astype(self.dtype)) return pycuda.gpuarray.to_gpu(_randn) def _inner_prod(self, x, y): return pycuda.gpuarray.dot(x.ravel().conj(), y.ravel()).real def _covector(self, x): return x.conj() VSpace.register(GPUArray, lambda x: ComplexGPUArrayVSpace(x) if cumisc.iscomplextype(x.dtype) else GPUArrayVSpace(x)) # for type_ in (float, np.float64, np.float32, np.float16): # GPUArrayVSpace.register(type_) # for type_ in (complex, np.complex64, np.complex128): # ComplexGPUArrayVSpace.register(type_) ### GPU VJPS ### GPU_DIFFERENTIABLE_FUNCTIONS = ( # skcuda.linalg
1.0j * np.ones(self.shape, dtype=self.dtype)) def standard_basis(self): for idxs in np.ndindex(*self.shape): for v in [1.0, 1.0j]: vect = np.zeros(self.shape, dtype=self.dtype) vect[idxs] = v yield vect def randn(self): return ( np.array(np.random.randn(*self.shape)).astype(self.dtype) + 1.0j * np.array(np.random.randn(*self.shape)).astype(self.dtype)) def _inner_prod(self, x, y): return np.real(np.dot(np.conj(np.ravel(x)), np.ravel(y))) def _covector(self, x): return np.conj(x) VSpace.register( np.ndarray, lambda x: ComplexArrayVSpace(x) if np.iscomplexobj(x) else ArrayVSpace(x)) for type_ in [float, np.float64, np.float32, np.float16]: ArrayVSpace.register(type_) for type_ in [complex, np.complex64, np.complex128]: ComplexArrayVSpace.register(type_)
from autograd.extend import VSpace class SparseArrayVSpace(VSpace): def __init__(self, value): self.t = type(value) self.value = self.t(value) self.shape = value.shape self.dtype = value.dtype @property def size(self): return self.value.size @property def ndim(self): return len(self.shape) def randn(self): a = cp.sparse.random(m=self.shape[0], n=self.shape[1]) return self.t(a) def zeros(self): return self.t(self.shape) VSpace.register(cp.sparse.csc_matrix, lambda x: SparseArrayVSpace(x)) VSpace.register(cp.sparse.csr_matrix, lambda x: SparseArrayVSpace(x)) VSpace.register(cp.sparse.coo_matrix, lambda x: SparseArrayVSpace(x)) VSpace.register(cp.sparse.dia_matrix, lambda x: SparseArrayVSpace(x))
return ( np.ones(self.shape, dtype=self.dtype) + 1.0j * np.ones(self.shape, dtype=self.dtype)) def standard_basis(self): for idxs in np.ndindex(*self.shape): for v in [1.0, 1.0j]: vect = np.zeros(self.shape, dtype=self.dtype) vect[idxs] = v yield vect def randn(self): return ( np.array(np.random.randn(*self.shape)).astype(self.dtype) + 1.0j * np.array(np.random.randn(*self.shape)).astype(self.dtype)) def _inner_prod(self, x, y): return np.real(np.dot(np.conj(np.ravel(x)), np.ravel(y))) def _covector(self, x): return np.conj(x) VSpace.register(np.ndarray, lambda x: ComplexArrayVSpace(x) if np.iscomplexobj(x) else ArrayVSpace(x)) for type_ in [float, np.float64, np.float32, np.float16]: ArrayVSpace.register(type_) for type_ in [complex, np.complex64, np.complex128]: ComplexArrayVSpace.register(type_)
return self.value.size @property def ndim(self): return len(self.shape) def randn(self): a = scipy.sparse.random(m=self.shape[0], n=self.shape[1]) return self.t(a) def zeros(self): return self.t(self.shape) def _inner_prod(self, x, y): x_vec = x.reshape((-1, 1)) y_vec = y.reshape((-1, 1)) dot_prod = np.sum(x_vec.T.dot(y_vec)) return dot_prod # def __eq__(self, other): # issametype = type(self) == type(other) # issamevals = (self.value != other.value).nnz == 0 # return issametype # and issamevals VSpace.register(scipy.sparse.csc_matrix, lambda x: SparseArrayVSpace(x)) VSpace.register(scipy.sparse.csr_matrix, lambda x: SparseArrayVSpace(x)) VSpace.register(scipy.sparse.coo_matrix, lambda x: SparseArrayVSpace(x)) VSpace.register(scipy.sparse.dia_matrix, lambda x: SparseArrayVSpace(x)) VSpace.register(np.matrixlib.defmatrix.matrix, lambda x: VSpace(x))