Index

_

__call__() (transformer_engine.jax.flax.DenseGeneral method) (transformer_engine.jax.flax.LayerNorm method) (transformer_engine.jax.flax.LayerNormDenseGeneral method) (transformer_engine.jax.flax.LayerNormMLP method) (transformer_engine.jax.flax.MultiHeadAttention method) (transformer_engine.jax.flax.RelativePositionBiases method) (transformer_engine.jax.flax.TransformerLayer method)

C

checkpoint() (in module transformer_engine.pytorch)

D

DelayedScaling (class in transformer_engine.common.recipe)

DenseGeneral (class in transformer_engine.jax.flax) DotProductAttention (class in transformer_engine.pytorch)

E

extend_logical_axis_rules() (in module transformer_engine.jax.flax)

F

Format (class in transformer_engine.common.recipe) forward() (transformer_engine.pytorch.DotProductAttention method) (transformer_engine.pytorch.LayerNormLinear method) (transformer_engine.pytorch.LayerNormMLP method) (transformer_engine.pytorch.Linear method) (transformer_engine.pytorch.TransformerLayer method)

fp8_autocast() (in module transformer_engine.jax) (in module transformer_engine.pytorch)

L

LayerNorm (class in transformer_engine.jax.flax) (class in transformer_engine.pytorch) LayerNormDenseGeneral (class in transformer_engine.jax.flax)

LayerNormLinear (class in transformer_engine.pytorch) LayerNormMLP (class in transformer_engine.jax.flax) (class in transformer_engine.pytorch) Linear (class in transformer_engine.pytorch)

M

MajorShardingType (class in transformer_engine.jax)

MultiHeadAttention (class in transformer_engine.jax.flax)

N

NVTE_Bias_Type (C++ enum) NVTE_Bias_Type::NVTE_NO_BIAS (C++ enumerator) NVTE_Bias_Type::NVTE_POST_SCALE_BIAS (C++ enumerator) NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS (C++ enumerator) nvte_cast_transpose (C++ function) nvte_cast_transpose_dbias (C++ function) nvte_cast_transpose_dbias_dgelu (C++ function) nvte_create_tensor (C++ function) nvte_cublas_gemm (C++ function) nvte_destroy_tensor (C++ function) nvte_dgeglu (C++ function) nvte_dgeglu_cast_transpose (C++ function) nvte_fp8_dequantize (C++ function) nvte_fp8_quantize (C++ function) nvte_fp8_transpose_dbias (C++ function) nvte_fused_attn_bwd_kvpacked (C++ function) nvte_fused_attn_bwd_qkvpacked (C++ function) nvte_fused_attn_fwd_kvpacked (C++ function) nvte_fused_attn_fwd_qkvpacked (C++ function) nvte_geglu (C++ function) nvte_gelu (C++ function) nvte_layernorm1p_bwd (C++ function) nvte_layernorm1p_fwd (C++ function) nvte_layernorm_bwd (C++ function) nvte_layernorm_fwd (C++ function) NVTE_Mask_Type (C++ enum) NVTE_Mask_Type::NVTE_CAUSAL_MASK (C++ enumerator) NVTE_Mask_Type::NVTE_NO_MASK (C++ enumerator) NVTE_Mask_Type::NVTE_PADDING_MASK (C++ enumerator) nvte_multi_cast_transpose (C++ function) NVTE_QKV_Layout (C++ enum) NVTE_QKV_Layout::NVTE_KV_INTERLEAVED (C++ enumerator) NVTE_QKV_Layout::NVTE_NOT_INTERLEAVED (C++ enumerator)

NVTE_QKV_Layout::NVTE_QKV_INTERLEAVED (C++ enumerator) nvte_scaled_masked_softmax_backward (C++ function) nvte_scaled_masked_softmax_forward (C++ function) nvte_scaled_softmax_backward (C++ function) nvte_scaled_softmax_forward (C++ function) nvte_scaled_upper_triang_masked_softmax_backward (C++ function) nvte_scaled_upper_triang_masked_softmax_forward (C++ function) nvte_tensor_amax (C++ function) nvte_tensor_data (C++ function) nvte_tensor_pack_create (C++ function) nvte_tensor_pack_destroy (C++ function) nvte_tensor_scale (C++ function) nvte_tensor_scale_inv (C++ function) nvte_tensor_shape (C++ function) nvte_tensor_type (C++ function) nvte_transpose (C++ function) NVTEDType (C++ enum) NVTEDType::kNVTEBFloat16 (C++ enumerator) NVTEDType::kNVTEByte (C++ enumerator) NVTEDType::kNVTEFloat16 (C++ enumerator) NVTEDType::kNVTEFloat32 (C++ enumerator) NVTEDType::kNVTEFloat8E4M3 (C++ enumerator) NVTEDType::kNVTEFloat8E5M2 (C++ enumerator) NVTEDType::kNVTEInt32 (C++ enumerator) NVTEDType::kNVTEInt64 (C++ enumerator) NVTEDType::kNVTENumTypes (C++ enumerator) NVTEShape (C++ struct) NVTEShape::data (C++ member) NVTEShape::ndim (C++ member) NVTETensor (C++ type) NVTETensorPack (C++ struct) NVTETensorPack::MAX_SIZE (C++ member) NVTETensorPack::size (C++ member) NVTETensorPack::tensors (C++ member)

O

onnx_export() (in module transformer_engine.pytorch)

R

RelativePositionBiases (class in transformer_engine.jax.flax)

S

ShardingResource (class in transformer_engine.jax)

ShardingType (class in transformer_engine.jax)

T

transformer_engine (C++ type) transformer_engine::DType (C++ enum) transformer_engine::DType::kBFloat16 (C++ enumerator) transformer_engine::DType::kByte (C++ enumerator) transformer_engine::DType::kFloat16 (C++ enumerator) transformer_engine::DType::kFloat32 (C++ enumerator) transformer_engine::DType::kFloat8E4M3 (C++ enumerator) transformer_engine::DType::kFloat8E5M2 (C++ enumerator) transformer_engine::DType::kInt32 (C++ enumerator) transformer_engine::DType::kInt64 (C++ enumerator) transformer_engine::DType::kNumTypes (C++ enumerator) transformer_engine::TensorWrapper (C++ struct) transformer_engine::TensorWrapper::amax (C++ function)

transformer_engine::TensorWrapper::data (C++ function) transformer_engine::TensorWrapper::dptr (C++ function) transformer_engine::TensorWrapper::dtype (C++ function) transformer_engine::TensorWrapper::operator= (C++ function), [1] transformer_engine::TensorWrapper::scale (C++ function) transformer_engine::TensorWrapper::scale_inv (C++ function) transformer_engine::TensorWrapper::shape (C++ function) transformer_engine::TensorWrapper::tensor_ (C++ member) transformer_engine::TensorWrapper::TensorWrapper (C++ function), [1], [2], [3], [4] transformer_engine::TensorWrapper::~TensorWrapper (C++ function) TransformerLayer (class in transformer_engine.jax.flax) (class in transformer_engine.pytorch) TransformerLayerType (class in transformer_engine.jax)

U

update_collections() (in module transformer_engine.jax)

update_fp8_metas() (in module transformer_engine.jax)