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)
