def get_handle(): dev = device.get_device_id() if dev not in _handles: handle = cutensor.Handle() cutensor.init(handle) _handles[dev] = handle return _handles[dev]
def set_random_state(rs): """Sets the state of the random number generator for the current device. Args: state(RandomState): Random state to set for the current device. """ if not isinstance(rs, RandomState): raise TypeError('Random state must be an instance of RandomState. ' 'Actual: {}'.format(type(rs))) _random_states[device.get_device_id()] = rs
def fill_diagonal(a, val, wrap=False): """Fills the main diagonal of the given array of any dimensionality. For an array `a` with ``a.ndim > 2``, the diagonal is the list of locations with indices ``a[i, i, ..., i]`` all identical. This function modifies the input array in-place, it does not return a value. Args: a (cupy.ndarray): The array, at least 2-D. val (scalar): The value to be written on the diagonal. Its type must be compatible with that of the array a. wrap (bool): If specified, the diagonal is "wrapped" after N columns. This affects only tall matrices. Examples -------- >>> a = cupy.zeros((3, 3), int) >>> cupy.fill_diagonal(a, 5) >>> a array([[5, 0, 0], [0, 5, 0], [0, 0, 5]]) .. seealso:: :func:`numpy.fill_diagonal` """ # The followings are imported from the original numpy if a.ndim < 2: raise ValueError('array must be at least 2-d') end = a.size if a.ndim == 2: step = a.shape[1] + 1 if not wrap: end = a.shape[1] * a.shape[1] else: if not numpy.alltrue(numpy.diff(a.shape) == 0): raise ValueError('All dimensions of input must be of equal length') step = 1 + numpy.cumprod(a.shape[:-1]).sum() val = cupy.asarray(val, dtype=a.dtype) dev_id = device.get_device_id() for arr in [a, val]: if arr.data.device_id != dev_id: raise ValueError('Array device must be same as the current ' 'device: array device = %d while current = %d' % (arr.data.device_id, dev_id)) typename = _scalar.get_typename(a.dtype) fill_diagonal_kernel = _fill_diagonal_kernel(typename, a.ndim, val.ndim) size = end // step + 1 a_ind = _carray.Indexer(a.shape) val_ind = _carray.Indexer(val.shape) fill_diagonal_kernel.kernel.linear_launch( size, (a, a_ind, 0, end, step, val, val_ind))
def get_handle(): dev = device.get_device_id() if dev not in _handles: _handles[dev] = Descriptor(cutensor.init(), cutensor.destroy) return _handles[dev]
def get_handle(): dev = device.get_device_id() if dev not in _handles: _handles[dev] = cutensor.create() return _handles[dev]
def n_free_blocks(self): dev = device.get_device_id() return self._pools[dev].n_free_blocks()
def free_all_free(self): dev = device.get_device_id() self._pools[dev].free_all_free()
def malloc(self, size): dev = device.get_device_id() return self._pools[dev].malloc(size)