def __setitem__(self, key, value): sliced = self.__getitem__(key) # Use fill for scalar values if type(value) == np.float32 or type(value) == np.float64 or \ type(value) == float: sliced.fill(value) elif type(value) == np.int32 or type(value) == np.int64 or \ type(value) == int: sliced.fill(value) elif self.tensor.shape == () or np.prod(self.tensor.shape) == 1: sliced.fill(value) elif np.sum(self.tensor.strides) == 0: view = GPUArray((1, ), dtype=self.tensor.dtype) view.fill(value) else: # Convert to correct dtype if necessary if value.dtype != self.tensor.dtype: new_value = np.ndarray(self.tensor.shape, dtype=self.tensor.dtype) new_value[:] = value value = new_value # Reshape to satisfy pycuda if necessary if sliced.shape != value.shape: sliced = self.tensor.reshape(value.shape) if self.is_contiguous and self.strides_contiguous(value): sliced[:] = value elif type(value) == GPUArray: self.from_other(value, sliced) else: contig_tensor = GPUArray(value.shape, self.tensor.dtype) contig_tensor[:] = value self.from_other(contig_tensor, sliced)
def __setitem__(self, key, value): sliced = self.__getitem__(key) # Use fill for scalar values # convert value to numpy if type(value) == float: value = np.float64(value) elif type(value) == int: value = np.int64(value) elif isinstance(value, np.ndarray): # handle 0-d and 1-d conversion to scalar if value.shape == (): value = value[()] elif value.shape == (1, ): value = value[0] # flex: added astype to deal with GPUArray dtype int16 # FLEX TODO: assumed same behavior for all cases if type(value) in (np.int32, np.int64, int, np.uint32, np.float32, np.float64): sliced.fill(value.astype(sliced.dtype)) elif self.tensor.shape == () or np.prod(self.tensor.shape) == 1: sliced.fill(value.astype(sliced.dtype)) elif np.sum(self.tensor.strides) == 0: view = GPUArray((1, ), dtype=self.tensor.dtype) view.fill(value.astype(sliced.dtype)) else: # Convert to correct dtype if necessary if value.dtype != self.tensor.dtype: new_value = np.ndarray(value.shape, dtype=self.tensor.dtype) new_value[:] = value value = new_value # Reshape to satisfy pycuda if necessary if sliced.shape != value.shape: sliced = self.tensor.reshape(value.shape) if self.is_contiguous and self.strides_contiguous(value): if sliced.shape == (): sliced.reshape((1, ))[:] = value.reshape((1, )) else: sliced[:] = value elif type(value) == GPUArray: self.from_other(value, sliced) else: contig_tensor = GPUArray(value.shape, self.tensor.dtype) contig_tensor[:] = value self.from_other(contig_tensor, sliced)