pub struct AbstractSortedSegmentOp<T, SIndex, Context, Reducer, const SparseFused: bool, InputAccessor> { /* private fields */ }Expand description
| @brief Segment reduction op with optional fused
| embedding lookup
|
| Base implementation for SortedSegmentXXX and
| SparseSortedSegmentXXX depending on SparseFused
| static argument.
|
| Inputs:
| 0: DATA - input embedding to do lookups in
| 1..P: AUX_ARG_ - optional additional arguments to be passed to the
| reducer, should have the same first dimension as
| SEGMENT_IDS (e.g. scalars in WeightedSum)
| # if SparseFused == true:
| P+1: INDICES - 1-D vector with indices to look up in DATA. Should have the
| same dimension as SEGMENT_IDS
| # P+1 if SparseFused == false:
| P+1 or P+2: SEGMENT_IDS - sorted segment ids 1-D vector
|
| Output:
|
| Tensor with first dimension of K, where K is
| the max segment id + 1. Rest of dimensions are
| decided by reducer but usually are the same
| size as extra dimensions of DATA
|
| bool SparseFused = true,
| class InputAccessor = BaseInputAccessor
Implementations§
source§impl<T, SIndex, Context, R: Reducer, const SparseFused: bool, InputAccessor> AbstractSortedSegmentOp<T, SIndex, Context, R, SparseFused, InputAccessor>
impl<T, SIndex, Context, R: Reducer, const SparseFused: bool, InputAccessor> AbstractSortedSegmentOp<T, SIndex, Context, R, SparseFused, InputAccessor>
pub fn run_on_device(&mut self) -> bool
pub fn do_run_with_type<IndexType>(&mut self) -> bool
pub fn do_run_with_value<IndexType, const FixedSize: i32>(&mut self) -> bool
Auto Trait Implementations§
impl<T, SIndex, Context, Reducer, const SparseFused: bool, InputAccessor> !RefUnwindSafe for AbstractSortedSegmentOp<T, SIndex, Context, Reducer, SparseFused, InputAccessor>
impl<T, SIndex, Context, Reducer, const SparseFused: bool, InputAccessor> !Send for AbstractSortedSegmentOp<T, SIndex, Context, Reducer, SparseFused, InputAccessor>
impl<T, SIndex, Context, Reducer, const SparseFused: bool, InputAccessor> !Sync for AbstractSortedSegmentOp<T, SIndex, Context, Reducer, SparseFused, InputAccessor>
impl<T, SIndex, Context, Reducer, const SparseFused: bool, InputAccessor> Unpin for AbstractSortedSegmentOp<T, SIndex, Context, Reducer, SparseFused, InputAccessor>where Context: Unpin, InputAccessor: Unpin, Reducer: Unpin, SIndex: Unpin, T: Unpin,
impl<T, SIndex, Context, Reducer, const SparseFused: bool, InputAccessor> !UnwindSafe for AbstractSortedSegmentOp<T, SIndex, Context, Reducer, SparseFused, InputAccessor>
Blanket Implementations§
§impl<T> Pointable for T
impl<T> Pointable for T
§impl<SS, SP> SupersetOf<SS> for SPwhere
SS: SubsetOf<SP>,
impl<SS, SP> SupersetOf<SS> for SPwhere SS: SubsetOf<SP>,
§fn to_subset(&self) -> Option<SS>
fn to_subset(&self) -> Option<SS>
self from the equivalent element of its
superset. Read more§fn is_in_subset(&self) -> bool
fn is_in_subset(&self) -> bool
self is actually part of its subset T (and can be converted to it).§fn to_subset_unchecked(&self) -> SS
fn to_subset_unchecked(&self) -> SS
self.to_subset but without any property checks. Always succeeds.§fn from_subset(element: &SS) -> SP
fn from_subset(element: &SS) -> SP
self to the equivalent element of its superset.