def compare_with_tensorflow(device_type, in_shape, axis, k, data_type, sorted): assert device_type in ["gpu", "cpu"] assert data_type in ["float32", "double", "int8", "int32", "int64"] flow.clear_default_session() func_config = flow.FunctionConfig() func_config.default_logical_view(flow.scope.mirrored_view()) func_config.default_data_type(flow.float) @flow.global_function(function_config=func_config) def TopKJob(input: oft.ListNumpy.Placeholder( tuple([dim + 10 for dim in in_shape]), dtype=type_name_to_flow_type[data_type], )): with flow.scope.placement(device_type, "0:0"): return flow.math.top_k(input, axis, k, sorted) input = (np.random.random(in_shape) * 100).astype( type_name_to_np_type[data_type]) # OneFlow of_out = TopKJob([input]).get().numpy_list()[0] # TensorFlow if k <= in_shape[axis]: perm = get_perm_when_transpose_axis_to_last_dim(len(in_shape), axis) x = tf.transpose(input, perm) _, indices = tf.math.top_k(x, k, sorted) tf_out = tf.transpose(indices, get_inversed_perm(perm)) else: tf_out = tf.argsort(input, axis, direction="DESCENDING", stable=True) assert np.array_equal(of_out, tf_out.numpy())
def forward(self, input): if self.dim == None: self.dim = -1 num_axes = len(input.shape) axis = self.dim if self.dim >= 0 else self.dim + num_axes assert 0 <= axis < num_axes, "axis out of range" if axis == num_axes - 1: if self.largest: indices = self._op_topk_last_dim(input)[0] else: neg_input = flow.experimental.mul(input, -1) indices = self._op_topk_last_dim(neg_input)[0] return (flow.experimental.gather(input, indices, dim=axis), indices) else: perm = get_perm_when_transpose_axis_to_last_dim(num_axes, axis) x = flow.F.transpose(input, perm=perm) if self.largest: indices = self._op_topk_last_dim(x)[0] else: neg_input = flow.experimental.mul(x, -1) indices = self._op_topk_last_dim(neg_input)[0] indices = flow.F.transpose(indices, perm=get_inversed_perm(perm)) return (flow.experimental.gather(input, indices, dim=axis), indices)
def sort( input: remote_blob_util.BlobDef, axis: int = -1, direction: str = "ASCENDING", name: Optional[str] = None, ) -> remote_blob_util.BlobDef: """This operator sorts the input Blob at specified axis. Args: input (remote_blob_util.BlobDef): A Blob axis (int, optional): dimension to be sorted. Defaults to the last dim (-1) direction (str, optional): The direction in which to sort the Blob values. If the direction is "ASCENDING", The order of input will be sorted as ascending, else, the order of input will be sorted as descending. Defaults to "ASCENDING". name (Optional[str], optional): The name for the operation. Defaults to None. Returns: remote_blob_util.BlobDef: The sorted Blob For example: .. code-block:: python import oneflow as flow import numpy as np import oneflow.typing as tp @flow.global_function() def sort_Job(x: tp.Numpy.Placeholder((5, )) ) -> tp.Numpy: return flow.sort(input=x) x = np.array([10, 2, 9, 3, 7]).astype("float32") out = sort_Job(x) # out [ 2. 3. 7. 9. 10.] """ assert direction in ["ASCENDING", "DESCENDING"] name = name if name is not None else id_util.UniqueStr("Sort_") num_axes = len(input.shape) axis = axis if axis >= 0 else axis + num_axes assert 0 <= axis < num_axes, "axis out of range" if axis == num_axes - 1: return _sort_at_last_dim(input, direction, name) else: perm = get_perm_when_transpose_axis_to_last_dim(num_axes, axis) x = flow.transpose(input, perm, False, True, name + "_transpose") x = _sort_at_last_dim(x, direction, name) return flow.transpose( x, get_inversed_perm(perm), False, True, name + "_inverse_transpose" )
def forward(self, input): if self.dim == None: input = self._flatten(input)[0] self.dim = 0 num_axes = len(input.shape) axis = self.dim if self.dim >= 0 else self.dim + num_axes assert 0 <= axis < num_axes, "axis out of range" if axis == num_axes - 1: x = self._op_softmax_last_dim(input)[0] if self.keepdim == True: x = self._expand_op(x) return x else: perm = get_perm_when_transpose_axis_to_last_dim(num_axes, axis) x = flow.tmp.transpose(input, perm=perm) x = self._op_softmax_last_dim(x)[0] x = self._expand_op(x)[0] x = flow.tmp.transpose(x, perm=get_inversed_perm(perm)) if self.keepdim == False: x = flow.tmp.squeeze(x, axis=[axis]) return x
def forward(self, input): if self.dim == None: input = flow.F.flatten(input) self.dim = 0 num_axes = len(input.shape) axis = self.dim if self.dim >= 0 else self.dim + num_axes assert 0 <= axis < num_axes, "axis out of range" if axis == num_axes - 1: x = flow.F.argmax(input) if self.keepdim == True: x = flow.experimental.unsqueeze(x, -1) return x else: perm = get_perm_when_transpose_axis_to_last_dim(num_axes, axis) x = flow.F.transpose(input, perm=perm) x = flow.F.argmax(x) x = flow.experimental.unsqueeze(x, -1) x = flow.F.transpose(x, perm=get_inversed_perm(perm)) if self.keepdim == False: x = x.squeeze(dim=[axis]) return x