Enum EmbeddingKind 

#[non_exhaustive]
#[repr(u16)]
pub enum EmbeddingKind {
    Embedding = 0,
    EmbeddingBackward = 1,
    EmbeddingBagSum = 2,
    EmbeddingBagMean = 3,
    EmbeddingBagSumBackward = 4,
    EmbeddingBagMeanBackward = 5,
    EmbeddingBagMax = 6,
    EmbeddingBagMaxBackward = 7,
}

Embedding-family op discriminant — Category M from the comprehensive plan.

Stored as u16 in crate::KernelSku::op when category == OpCategory::Embedding. Phase 7 Milestone 7.5 wires:

• Self::Embedding (FW + BW): row-lookup out[i, :] = weight[indices[i], :] with optional padding_idx that emits an all-zero row at FW and skips accumulation at BW.
• Self::EmbeddingBagSum / Self::EmbeddingBagMean (FW + BW): bag-reduced row lookup — out[b, :] = reduce(weight[indices[k], :] for k in offsets[b]..offsets[b+1]). Mode determines the reducer (sum / divide-by-bag-size). EmbeddingBagMax is deferred (needs argmax tracking for BW).

Index dtype is i32 only (i64 deferred). FW kernels emit f32, f64, f16, bf16 (pure copy / reduce); BW kernels emit f32, f64 (atomicAdd).

Variants (Non-exhaustive)

Non-exhaustive enums could have additional variants added in future. Therefore, when matching against variants of non-exhaustive enums, an extra wildcard arm must be added to account for any future variants.

Embedding = 0

embedding(weight, indices, padding_idx) — out[i, :] = weight[indices[i], :]. PyTorch torch.nn.functional.embedding.

EmbeddingBackward = 1

Gradient of Self::Embedding: dweight[indices[i], :] += dout[i, :] (atomicAdd), skipping rows where indices[i] == padding_idx.

EmbeddingBagSum = 2

embedding_bag(weight, indices, offsets, mode=Sum). PyTorch torch.nn.functional.embedding_bag with mode='sum'.

EmbeddingBagMean = 3

embedding_bag(weight, indices, offsets, mode=Mean). PyTorch torch.nn.functional.embedding_bag with mode='mean'.

EmbeddingBagSumBackward = 4

Gradient of embedding_bag (Sum-mode): dweight[indices[k], :] += dout[b, :] for k in bag b (atomicAdd).

EmbeddingBagMeanBackward = 5

Gradient of embedding_bag (Mean-mode): dweight[indices[k], :] += dout[b, :] / bag_size(b) (atomicAdd).

EmbeddingBagMax = 6

embedding_bag(weight, indices, offsets, mode=Max) — reserved. Max-mode requires argmax tracking on FW (the per-feature index of the contributing row) so the BW can scatter into just that row — different plan shape; deferred.

EmbeddingBagMaxBackward = 7

Gradient of embedding_bag (Max-mode) — reserved.

Trait Implementations

impl Clone for EmbeddingKind

fn clone(&self) -> EmbeddingKind

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Copy for EmbeddingKind

impl Debug for EmbeddingKind

fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error>

Formats the value using the given formatter.

impl Eq for EmbeddingKind

impl Hash for EmbeddingKind

fn hash<__H>(&self, state: &mut __H)

where __H: Hasher,

Feeds this value into the given Hasher.

fn hash_slice<H>(data: &[Self], state: &mut H)

where H: Hasher, Self: Sized,

Feeds a slice of this type into the given Hasher.

impl PartialEq for EmbeddingKind

fn eq(&self, other: &EmbeddingKind) -> bool

Tests for self and other values to be equal, and is used by ==.

fn ne(&self, other: &Rhs) -> bool

Tests for !=. The default implementation is almost always sufficient, and should not be overridden without very good reason.

impl StructuralPartialEq for EmbeddingKind