Mojo package

nn

Provides neural network operators for deep learning models.

Modules

• ​activations: The module contains implementations of activation functions.
• ​arange:
• ​arg_nonzero:
• ​argmaxmin:
• ​argmaxmin_gpu:
• ​argsort:
• ​bicubic: This module provides CPU and GPU implementations for bicubic interpolation.
• ​broadcast:
• ​concat:
• ​conv:
• ​conv_transpose:
• ​conv_utils:
• ​cumsum:
• ​flash_attention:
• ​fold: Implements the fold operation.
• ​fused_qk_rope:
• ​gather_scatter:
• ​image:
• ​index_tensor:
• ​irfft: Inverse real FFT kernel using cuFFT.
• ​kv_cache:
• ​kv_cache_ragged:
• ​mha:
• ​mha_amd:
• ​mha_cross:
• ​mha_fa3_utils:
• ​mha_mask:
• ​mha_operand:
• ​mha_score_mod:
• ​mha_sm100:
• ​mha_sm90:
• ​mha_tile_scheduler:
• ​mha_utils:
• ​mla:
• ​moe:
• ​nms:
• ​normalization:
• ​pad:
• ​pad_gpu:
• ​pool:
• ​rand_normal:
• ​rand_uniform:
• ​randn:
• ​repeat_interleave:
• ​reshape:
• ​resize:
• ​roi_align:
• ​rope:
• ​sampling:
• ​shapes:
• ​slice:
• ​softmax:
• ​split:
• ​tile:
• ​topk:
• ​toppminp:
• ​toppminp_gpu:

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit