All modules for which code is available

• _hashlib
• abc
• aitemplate.backend.build_cache_base
• aitemplate.backend.builder
• aitemplate.backend.codegen
• aitemplate.backend.cuda.target_def
• aitemplate.backend.profiler_cache
• aitemplate.backend.profiler_runner
• aitemplate.backend.registry
• aitemplate.backend.rocm.target_def
• aitemplate.backend.target
• aitemplate.backend.task_runner
• aitemplate.compiler.base
• aitemplate.compiler.dtype
• aitemplate.compiler.model
• aitemplate.compiler.ops.attention.flash_attention
• aitemplate.compiler.ops.attention.mem_eff_attention
• aitemplate.compiler.ops.b2b_bmm.classic_b2b_bmm
• aitemplate.compiler.ops.b2b_bmm.fmha_style_b2b_bmm
• aitemplate.compiler.ops.b2b_bmm.grouped_classic_b2b_bmm
• aitemplate.compiler.ops.b2b_bmm.grouped_fmha_style_b2b_bmm
• aitemplate.compiler.ops.common.elementwise
• aitemplate.compiler.ops.common.epilogue
• aitemplate.compiler.ops.common.fused_elementwise
• aitemplate.compiler.ops.common.int_elementwise
• aitemplate.compiler.ops.common.python_ops
• aitemplate.compiler.ops.common.view_ops
• aitemplate.compiler.ops.conv.conv2d
• aitemplate.compiler.ops.conv.conv2d_bias
• aitemplate.compiler.ops.conv.conv2d_bias_add
• aitemplate.compiler.ops.conv.conv2d_bias_add_hardswish
• aitemplate.compiler.ops.conv.conv2d_bias_add_relu
• aitemplate.compiler.ops.conv.conv2d_bias_few_channels
• aitemplate.compiler.ops.conv.conv2d_bias_hardswish
• aitemplate.compiler.ops.conv.conv2d_bias_hardswish_few_channels
• aitemplate.compiler.ops.conv.conv2d_bias_relu
• aitemplate.compiler.ops.conv.conv2d_bias_relu_few_channels
• aitemplate.compiler.ops.conv.conv2d_bias_sigmoid
• aitemplate.compiler.ops.conv.conv2d_depthwise
• aitemplate.compiler.ops.conv.conv2d_depthwise_bias
• aitemplate.compiler.ops.conv.conv3d
• aitemplate.compiler.ops.conv.conv3d_bias
• aitemplate.compiler.ops.conv.depthwise_conv3d
• aitemplate.compiler.ops.conv.transposed_conv2d
• aitemplate.compiler.ops.conv.transposed_conv2d_bias
• aitemplate.compiler.ops.conv.transposed_conv2d_bias_relu
• aitemplate.compiler.ops.embedding.bert_embeddings
• aitemplate.compiler.ops.gemm_epilogue_vistor.bmm_rcr_softmax
• aitemplate.compiler.ops.gemm_epilogue_vistor.dual_bmm_rrr_div
• aitemplate.compiler.ops.gemm_epilogue_vistor.dual_gemm_rcr_fast_gelu
• aitemplate.compiler.ops.gemm_epilogue_vistor.dual_gemm_rcr_silu
• aitemplate.compiler.ops.gemm_epilogue_vistor.gemm_rcr_bias_softmax
• aitemplate.compiler.ops.gemm_epilogue_vistor.gemm_rcr_softmax
• aitemplate.compiler.ops.gemm_special.batched_dense_vec_jagged_2d_mul
• aitemplate.compiler.ops.gemm_special.bmm_rcr_n1
• aitemplate.compiler.ops.gemm_special.bmm_rrr_k1_tanh
• aitemplate.compiler.ops.gemm_special.gemm_rrr_small_nk
• aitemplate.compiler.ops.gemm_universal.bmm_rcr_permute
• aitemplate.compiler.ops.gemm_universal.bmm_rrr_permute
• aitemplate.compiler.ops.gemm_universal.bmm_softmax_bmm
• aitemplate.compiler.ops.gemm_universal.bmm_softmax_bmm_permute
• aitemplate.compiler.ops.gemm_universal.bmm_xxx
• aitemplate.compiler.ops.gemm_universal.bmm_xxx_add
• aitemplate.compiler.ops.gemm_universal.gemm_common
• aitemplate.compiler.ops.gemm_universal.gemm_rcr
• aitemplate.compiler.ops.gemm_universal.gemm_rcr_bias
• aitemplate.compiler.ops.gemm_universal.gemm_rcr_bias_add
• aitemplate.compiler.ops.gemm_universal.gemm_rcr_bias_add_add
• aitemplate.compiler.ops.gemm_universal.gemm_rcr_bias_add_add_relu
• aitemplate.compiler.ops.gemm_universal.gemm_rcr_bias_add_relu
• aitemplate.compiler.ops.gemm_universal.gemm_rcr_bias_fast_gelu
• aitemplate.compiler.ops.gemm_universal.gemm_rcr_bias_gelu
• aitemplate.compiler.ops.gemm_universal.gemm_rcr_bias_hardswish
• aitemplate.compiler.ops.gemm_universal.gemm_rcr_bias_mul
• aitemplate.compiler.ops.gemm_universal.gemm_rcr_bias_mul_add
• aitemplate.compiler.ops.gemm_universal.gemm_rcr_bias_mul_tanh
• aitemplate.compiler.ops.gemm_universal.gemm_rcr_bias_permute
• aitemplate.compiler.ops.gemm_universal.gemm_rcr_bias_relu
• aitemplate.compiler.ops.gemm_universal.gemm_rcr_bias_sigmoid
• aitemplate.compiler.ops.gemm_universal.gemm_rcr_bias_sigmoid_mul
• aitemplate.compiler.ops.gemm_universal.gemm_rcr_bias_sigmoid_mul_tanh
• aitemplate.compiler.ops.gemm_universal.gemm_rcr_bias_swish
• aitemplate.compiler.ops.gemm_universal.gemm_rcr_bias_tanh
• aitemplate.compiler.ops.gemm_universal.gemm_rcr_fast_gelu
• aitemplate.compiler.ops.gemm_universal.gemm_rcr_permute
• aitemplate.compiler.ops.gemm_universal.gemm_rcr_permute_elup1
• aitemplate.compiler.ops.gemm_universal.gemm_rrr
• aitemplate.compiler.ops.gemm_universal.gemm_rrr_bias
• aitemplate.compiler.ops.gemm_universal.gemm_rrr_bias_permute
• aitemplate.compiler.ops.gemm_universal.gemm_rrr_permute
• aitemplate.compiler.ops.gemm_universal.group_gemm_rcr
• aitemplate.compiler.ops.gemm_universal.group_gemm_rcr_bias
• aitemplate.compiler.ops.gemm_universal.group_gemm_rcr_bias_relu
• aitemplate.compiler.ops.gemm_universal.group_gemm_rcr_bias_sigmoid
• aitemplate.compiler.ops.gemm_universal.perm021fc_ccr
• aitemplate.compiler.ops.gemm_universal.perm021fc_ccr_bias
• aitemplate.compiler.ops.gemm_universal.perm021fc_ccr_bias_permute
• aitemplate.compiler.ops.gemm_universal.perm021fc_crc
• aitemplate.compiler.ops.gemm_universal.perm021fc_crc_bias
• aitemplate.compiler.ops.gemm_universal.perm102_bmm_rcr
• aitemplate.compiler.ops.gemm_universal.perm102_bmm_rcr_bias
• aitemplate.compiler.ops.gemm_universal.perm102_bmm_rrr
• aitemplate.compiler.ops.gemm_universal.perm102_bmm_rrr_bias
• aitemplate.compiler.ops.groupnorm.groupnorm
• aitemplate.compiler.ops.groupnorm.groupnorm_swish
• aitemplate.compiler.ops.jagged.jagged_lengths_to_offsets
• aitemplate.compiler.ops.jagged.jagged_lengths_to_presences
• aitemplate.compiler.ops.layernorm.batch_layernorm_sigmoid_mul
• aitemplate.compiler.ops.layernorm.group_layernorm
• aitemplate.compiler.ops.layernorm.group_layernorm_sigmoid_mul
• aitemplate.compiler.ops.layernorm.layernorm
• aitemplate.compiler.ops.layernorm.layernorm_sigmoid_mul
• aitemplate.compiler.ops.padding.ndhwc3to8
• aitemplate.compiler.ops.padding.nhwc3to4
• aitemplate.compiler.ops.padding.nhwc3to8
• aitemplate.compiler.ops.padding.pad_last_dim
• aitemplate.compiler.ops.pool.avg_pool2d
• aitemplate.compiler.ops.pool.max_pool2d
• aitemplate.compiler.ops.reduce.reduce_max
• aitemplate.compiler.ops.reduce.reduce_mean
• aitemplate.compiler.ops.reduce.reduce_min
• aitemplate.compiler.ops.reduce.reduce_sum
• aitemplate.compiler.ops.reduce.var
• aitemplate.compiler.ops.reduce.vector_norm
• aitemplate.compiler.ops.softmax.softmax
• aitemplate.compiler.ops.tensor.argmax
• aitemplate.compiler.ops.tensor.batch_gather
• aitemplate.compiler.ops.tensor.cast
• aitemplate.compiler.ops.tensor.chunk
• aitemplate.compiler.ops.tensor.concatenate
• aitemplate.compiler.ops.tensor.concatenate_tanh
• aitemplate.compiler.ops.tensor.dynamic_slice
• aitemplate.compiler.ops.tensor.expand
• aitemplate.compiler.ops.tensor.full
• aitemplate.compiler.ops.tensor.gather
• aitemplate.compiler.ops.tensor.identity
• aitemplate.compiler.ops.tensor.index_select
• aitemplate.compiler.ops.tensor.jagged_to_padded_dense
• aitemplate.compiler.ops.tensor.masked_select
• aitemplate.compiler.ops.tensor.padded_dense_to_jagged
• aitemplate.compiler.ops.tensor.permute
• aitemplate.compiler.ops.tensor.permute021
• aitemplate.compiler.ops.tensor.permute0213
• aitemplate.compiler.ops.tensor.permute102
• aitemplate.compiler.ops.tensor.permute210
• aitemplate.compiler.ops.tensor.relational
• aitemplate.compiler.ops.tensor.size
• aitemplate.compiler.ops.tensor.slice_reshape_scatter
• aitemplate.compiler.ops.tensor.slice_scatter
• aitemplate.compiler.ops.tensor.split
• aitemplate.compiler.ops.tensor.topk
• aitemplate.compiler.ops.tensor.transpose
• aitemplate.compiler.ops.tensor.where
• aitemplate.compiler.ops.upsample.upsampling2d
• aitemplate.compiler.ops.upsample.upsampling2d_add
• aitemplate.compiler.ops.vision_ops.nms.batched_nms
• aitemplate.compiler.ops.vision_ops.nms.efficient_nms
• aitemplate.compiler.ops.vision_ops.nms.nms
• aitemplate.compiler.ops.vision_ops.roi_ops.multi_level_roi_align
• aitemplate.compiler.ops.vision_ops.roi_ops.roi_align
• aitemplate.compiler.stable_set
• aitemplate.compiler.symbolic
• aitemplate.compiler.tensor_accessor
• aitemplate.compiler.transform.apply_padding
• aitemplate.compiler.transform.bind_constants
• aitemplate.compiler.transform.constant_folding
• aitemplate.compiler.transform.dedup_make_jagged_ops
• aitemplate.compiler.transform.fuse_bmm_permute
• aitemplate.compiler.transform.fuse_conv_elementwise
• aitemplate.compiler.transform.fuse_conv_patterns
• aitemplate.compiler.transform.fuse_duplicate_fused_elementwise
• aitemplate.compiler.transform.fuse_expand_bmm
• aitemplate.compiler.transform.fuse_group_ops
• aitemplate.compiler.transform.fuse_mm_elementwise
• aitemplate.compiler.transform.fuse_mm_reshape_permute
• aitemplate.compiler.transform.fuse_ops
• aitemplate.compiler.transform.fuse_parallel_gemms
• aitemplate.compiler.transform.fuse_permute_bmm_and_gemm
• aitemplate.compiler.transform.mark_param_tensor
• aitemplate.compiler.transform.memory_planning
• aitemplate.compiler.transform.move_view_ops
• aitemplate.compiler.transform.name_graph
• aitemplate.compiler.transform.optimize_graph
• aitemplate.compiler.transform.profile
• aitemplate.compiler.transform.refine_graph
• aitemplate.compiler.transform.remove_elementwise_no_ops
• aitemplate.compiler.transform.remove_no_ops
• aitemplate.compiler.transform.remove_unused_ops
• aitemplate.compiler.transform.split_large_concat_ops
• aitemplate.compiler.transform.split_large_slice_scatter_ops
• aitemplate.compiler.transform.split_large_split_ops
• aitemplate.compiler.transform.toposort
• aitemplate.compiler.transform.transform_memory_ops
• aitemplate.compiler.transform.transform_merge_view_ops
• aitemplate.compiler.transform.transform_odd_alignment
• aitemplate.compiler.transform.transform_permute_to_reshape
• aitemplate.compiler.transform.transform_special_ops
• aitemplate.compiler.transform.transform_strided_ops
• aitemplate.frontend.nn.attention
• aitemplate.frontend.nn.container
• aitemplate.frontend.nn.conv1d
• aitemplate.frontend.nn.conv2d.conv2d
• aitemplate.frontend.nn.conv2d.conv2d_bias
• aitemplate.frontend.nn.conv2d.conv2d_bias_add_hardswish
• aitemplate.frontend.nn.conv2d.conv2d_bias_add_relu
• aitemplate.frontend.nn.conv2d.conv2d_bias_few_channels
• aitemplate.frontend.nn.conv2d.conv2d_bias_hardswish
• aitemplate.frontend.nn.conv2d.conv2d_bias_hardswish_few_channels
• aitemplate.frontend.nn.conv2d.conv2d_bias_relu
• aitemplate.frontend.nn.conv2d.conv2d_bias_relu_few_channels
• aitemplate.frontend.nn.conv2d.conv2d_bias_sigmoid
• aitemplate.frontend.nn.conv2d.conv2d_depthwise
• aitemplate.frontend.nn.conv2d.conv2d_depthwise_bias
• aitemplate.frontend.nn.conv2d.transposed_conv2d_bias
• aitemplate.frontend.nn.conv2d.transposed_conv2d_bias_relu
• aitemplate.frontend.nn.conv3d
• aitemplate.frontend.nn.dropout
• aitemplate.frontend.nn.dual_gemm
• aitemplate.frontend.nn.embedding
• aitemplate.frontend.nn.fpn_proposal
• aitemplate.frontend.nn.group_norm
• aitemplate.frontend.nn.identity
• aitemplate.frontend.nn.layer_norm
• aitemplate.frontend.nn.linear
• aitemplate.frontend.nn.module
• aitemplate.frontend.nn.multiscale_attention
• aitemplate.frontend.nn.padding
• aitemplate.frontend.nn.pool2d
• aitemplate.frontend.nn.proposal
• aitemplate.frontend.nn.roi_ops
• aitemplate.frontend.nn.upsample
• aitemplate.frontend.nn.vanilla_attention
• aitemplate.frontend.nn.view_ops
• aitemplate.testing.benchmark_ait
• aitemplate.testing.benchmark_pt
• aitemplate.testing.detect_target
• aitemplate.utils.debug_settings
• aitemplate.utils.environ
• aitemplate.utils.graph_utils
• aitemplate.utils.misc
• aitemplate.utils.serialization.serdes_code
• aitemplate.utils.shape_utils
• aitemplate.utils.tensor_utils
• aitemplate.utils.visualization.plot
• collections
• copy
• dataclasses
• datetime
• enum
• numbers
• pathlib
• queue
• time