cuda.tile.argmax

cuda.tile.argmax(x, /, axis=None, *, keepdims=False)

Performs argmax reduction on tile along the axis.

Parameters:

• x (Tile) – input tile.

• axis (None | const int | tuple[const int,...]) – the axis for reduction. The default, axis=None, will reduce all of the elements. For argmin and argmax, tuple of axis is not supported.

• keepdims (const bool) – If true, preserves the number of dimension from the input tile.

Returns:

Examples

Reduce all axes.

SnippetComplete Example

tx = ct.arange(8, dtype=ct.int32).reshape((2, 4))
print("input:", tx)
print("reduced:", ct.argmax(tx, None))

Copy

import cuda.tile as ct
import torch

@ct.kernel
def kernel():
    tx = ct.arange(8, dtype=ct.int32).reshape((2, 4))
    print("input:", tx)
    print("reduced:", ct.argmax(tx, None))


torch.cuda.init()
ct.launch(torch.cuda.current_stream(), (1,), kernel, ())
torch.cuda.synchronize()

Output

input: [[0, 1, 2, 3], [4, 5, 6, 7]]
reduced: 7

Reduce axis 1 and keepdims.

SnippetComplete Example

tx = ct.arange(8, dtype=ct.int32).reshape((2, 4))
print("input:", tx)
print("reduced:", ct.argmax(tx, 1, keepdims=True))

Copy

import cuda.tile as ct
import torch

@ct.kernel
def kernel():
    tx = ct.arange(8, dtype=ct.int32).reshape((2, 4))
    print("input:", tx)
    print("reduced:", ct.argmax(tx, 1, keepdims=True))


torch.cuda.init()
ct.launch(torch.cuda.current_stream(), (1,), kernel, ())
torch.cuda.synchronize()

Output

input: [[0, 1, 2, 3], [4, 5, 6, 7]]
reduced: [[3], [3]]

If there are ties, the smallest index is returned.

SnippetComplete Example

tx = ct.zeros(4, dtype=ct.int32)
tx = ct.cat((tx, tx + 1), axis=0)
print("input:", tx)
print("reduced:", ct.argmax(tx, 0))

Copy

import cuda.tile as ct
import torch

@ct.kernel
def kernel():
    tx = ct.zeros(4, dtype=ct.int32)
    tx = ct.cat((tx, tx + 1), axis=0)
    print("input:", tx)
    print("reduced:", ct.argmax(tx, 0))


torch.cuda.init()
ct.launch(torch.cuda.current_stream(), (1,), kernel, ())
torch.cuda.synchronize()

Output

input: [0, 0, 0, 0, 1, 1, 1, 1]
reduced: 4

Reduce over tuple of axes is not supported for argmax.

Return type:

Tile