Python API
- aitemplate.compiler
- aitemplate.compiler.ops
EpilogueOpFuncEnumJaggedDimJaggedIntVarargmaxavg_pool2dbatch_gatherbatch_layernorm_sigmoid_mulbatched_dense_vec_jagged_2d_mulbatched_nmsbert_embeddingsbmm_cccbmm_ccc_addbmm_ccrbmm_ccr_addbmm_crcbmm_crc_addbmm_crrbmm_crr_addbmm_rccbmm_rcc_addbmm_rcrbmm_rcr_addbmm_rcr_n1bmm_rcr_permutebmm_rcr_softmaxbmm_rrcbmm_rrc_addbmm_rrrbmm_rrr_addbmm_rrr_k1_tanhbmm_rrr_permutebmm_softmax_bmmbmm_softmax_bmm_permutecastchunkclampclassic_b2b_bmmconcatenateconcatenate_tanhconv2dconv2d_biasconv2d_bias_addconv2d_bias_add_hardswishconv2d_bias_add_reluconv2d_bias_few_channelsconv2d_bias_hardswishconv2d_bias_hardswish_few_channelsconv2d_bias_reluconv2d_bias_relu_few_channelsconv2d_bias_sigmoidconv2d_depthwiseconv2d_depthwise_biasconv3dconv3d_biasconvert_shape_to_IntVar()depthwise_conv3ddual_bmm_rrr_divdual_gemm_rcr_fast_geludual_gemm_rcr_siludynamic_sliceefficient_nmselementwiseeqexpandflash_attentionflattenfmha_style_b2b_bmmfullfused_elementwisegathergegemm_rcrgemm_rcr_biasgemm_rcr_bias_addgemm_rcr_bias_add_addgemm_rcr_bias_add_add_relugemm_rcr_bias_add_relugemm_rcr_bias_fast_gelugemm_rcr_bias_gelugemm_rcr_bias_hardswishgemm_rcr_bias_mulgemm_rcr_bias_mul_addgemm_rcr_bias_mul_tanhgemm_rcr_bias_permutegemm_rcr_bias_relugemm_rcr_bias_sigmoidgemm_rcr_bias_sigmoid_mulgemm_rcr_bias_sigmoid_mul_tanhgemm_rcr_bias_softmaxgemm_rcr_bias_swishgemm_rcr_bias_tanhgemm_rcr_fast_gelugemm_rcr_permutegemm_rcr_permute_elup1gemm_rcr_softmaxgemm_rrrgemm_rrr_biasgemm_rrr_bias_permutegemm_rrr_permutegemm_rrr_small_nkgen_int_var_min_max()getitemgroup_gemm_rcrgroup_gemm_rcr_biasgroup_gemm_rcr_bias_relugroup_gemm_rcr_bias_sigmoidgroup_layernormgroup_layernorm_sigmoid_mulgroup_normgroup_norm_swishgrouped_classic_b2b_bmmgrouped_fmha_style_b2b_bmmgtidentityindex_selectint_elementwiseis_symbolic()jagged_lengths_to_offsetsjagged_lengths_to_presencesjagged_to_padded_denselayernormlayernorm_sigmoid_mullelist_constructltmake_jaggedmasked_selectmax_pool2dmem_eff_attentionmulti_level_roi_alignndhwc3to8nenhwc3to4nhwc3to8nmsnormalize_dtype()pad_last_dimpadded_dense_to_jaggedperm021fc_ccrperm021fc_ccr_biasperm021fc_ccr_bias_permuteperm021fc_crcperm021fc_crc_biasperm102_bmm_rcrperm102_bmm_rcr_biasperm102_bmm_rrrperm102_bmm_rrr_biaspermutepermute021permute0213permute102permute210reduce_maxreduce_meanreduce_minreduce_sumreshaperoi_alignsimplify_intvar_values()sizeslice_reshape_scatterslice_scattersoftmaxsplitsqueezetopktransposetransposed_conv2dtransposed_conv2d_biastransposed_conv2d_bias_relutuple_constructunsqueezeupsampling2dupsampling2d_addvarvector_normwhere
- aitemplate.compiler.transform
- apply_padding
- bind_constants
- constant_folding
- fuse_conv_elementwise
- fuse_group_ops
- fuse_mm_elementwise
- fuse_ops
- fuse_parallel_gemms
- fuse_permute_bmm
- fuse_split
- mark_param_tensor
- memory_planning
- name_graph
- optimize_graph
- profile
- refine_graph
- remove_no_ops
- remove_unused_ops
- toposort
- transform_memory_ops
- transform_odd_alignment
- transform_special_ops
- transform_strided_op_and_view_op
- transform_strided_ops
- transform_strided_slice
- aitemplate.backend
- aitemplate.backend.cuda
- aitemplate.backend.rocm
- aitemplate.frontend
AvgPool2dBertEmbeddingsConv1dConv2dConv2dBiasConv2dBiasAddHardswishConv2dBiasAddReluConv2dBiasFewChannelsConv2dBiasHardswishConv2dBiasHardswishFewChannelsConv2dBiasReluConv2dBiasReluFewChannelsConv2dBiasSigmoidConv2dDepthwiseConv2dDepthwiseBiasConv3dConvTranspose2dBiasConvTranspose2dBiasReluCrossAttentionDropPathDropoutEmbeddingFPNProposalFPNRoiAlignFlashAttentionFlattenGroupNormIdentityLayerNormLinearMaxPool2dModuleModuleDictModuleListMultiScaleBlockMultiheadAttentionNdhwc3to8Nhwc3to8ProposalReshapeRoiAlignScaledDotProductAttentionSequentialT5DenseGatedGeluDenseTensorUpsampling2dUpsampling2dAddVanillaCrossAttentionVanillaMultiheadAttentionViewavg_pool2dconv3dconv3d_biasdepthwise_conv3ddetect_target()flattenmax_pool2dmulti_level_roi_alignndhwc3to8nhwc3to8reshaperoi_alignsqueezeunsqueezeupsampling2dupsampling2d_addvanilla_attention()
- aitemplate.testing
- aitemplate.utils