Mojo package

nn

Provides neural network operators for deep learning models.

Packages

• ​attention: Multi-head attention (MHA) and multi-head latent attention (MLA) kernels.
• ​conv: Convolution kernels (1D, 2D, 3D, transposed).

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:
• ​cumsum:
• ​fold: Implements the fold operation.
• ​fused_qk_rope:
• ​gather_scatter:
• ​gemv_partial_norm: Fused matvec (M=1) + partial RMS norm on B200.
• ​image:
• ​index_fp8:
• ​index_tensor:
• ​irfft: Inverse real FFT kernel using cuFFT.
• ​kv_cache:
• ​kv_cache_ragged:
• ​learnable_2d_interp_pos_emb: Learnable 2D interpolated position embedding (Kimi K2.5 MoonViT3d).
• ​moe:
• ​nms:
• ​normalization:
• ​pad:
• ​pad_gpu:
• ​pool:
• ​rand_normal:
• ​rand_uniform:
• ​randn:
• ​repeat_interleave:
• ​reshape:
• ​resize:
• ​roi_align:
• ​rope:
• ​rope_split_store: Fused rope + split + KV store kernel.
• ​sampling:
• ​shapes:
• ​shard_and_stack:
• ​slice:
• ​softmax:
• ​spatial_merge:
• ​split:
• ​tile:
• ​topk:
• ​topk_fi:
• ​toppminp:
• ​toppminp_gpu:
• ​tpool_patch_merger:

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