For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

arg_nonzero_shape

def arg_nonzero_shape[dtype: DType](input_buffer: TileTensor[dtype, Storage=input_buffer.Storage, address_space=input_buffer.address_space, linear_idx_type=input_buffer.linear_idx_type, element_size=input_buffer.element_size]) -> IndexList[Int(2)]

Return [NumNonZeros, InputRank] where NumNonZeros are the number of non-zero elements in the input.

Parameters:

dtype (DType): The element dtype.

Args:

input_buffer (TileTensor[dtype, Storage=input_buffer.Storage, address_space=input_buffer.address_space, linear_idx_type=input_buffer.linear_idx_type, element_size=input_buffer.element_size]): The tensor to count the non-zeros in.

Returns:

IndexList[Int(2)]: Shape of the arg_nonzero kernel for this input [NumNonZeros, InputRank].