Struct QkvBase

pub struct QkvBase<S = OwnedRepr<f64>, D = Ix2>
where
    D: Dimension,
    S: RawData,
{ /* private fields */ }

QkvBase is a container for the query, key, and value arrays used in the attention mechanism of the transformer model.

Implementations

impl<A, S, D> QkvBase<S, D>

where D: Dimension, S: RawData<Elem = A>,

pub fn builder<Sh, F>(shape: Sh, builder: F) -> Self

where F: Fn(D) -> ArrayBase<S, D>, Sh: ShapeBuilder<Dim = D>,

pub fn from_elem<Sh>(shape: Sh, value: A) -> Self

where Sh: ShapeBuilder<Dim = D>, A: Clone, S: DataOwned,

pub fn as_qkv( &self, ) -> (ArrayView<'_, A, D>, ArrayView<'_, A, D>, ArrayView<'_, A, D>)

where S: Data,

pub fn into_qkv(self) -> (ArrayBase<S, D>, ArrayBase<S, D>, ArrayBase<S, D>)

Consumes the store and returns a three-tuple consisting of the query, key, and value arrays respectively.

pub fn qkv(&self) -> (&ArrayBase<S, D>, &ArrayBase<S, D>, &ArrayBase<S, D>)

pub fn new<Sh: ShapeBuilder<Dim = D>>(shape: Sh) -> Self

where A: Default, S: DataOwned,

pub fn ones<Sh: ShapeBuilder<Dim = D>>(shape: Sh) -> Self

where A: Clone + One, S: DataOwned,

pub fn zeros<Sh: ShapeBuilder<Dim = D>>(shape: Sh) -> Self

where A: Clone + Zero, S: DataOwned,

pub fn q(&self) -> &ArrayBase<S, D>

pub fn q_mut(&mut self) -> &mut ArrayBase<S, D>

pub fn k(&self) -> &ArrayBase<S, D>

pub fn k_mut(&mut self) -> &mut ArrayBase<S, D>

pub fn v(&self) -> &ArrayBase<S, D>

pub fn v_mut(&mut self) -> &mut ArrayBase<S, D>

pub fn dim(&self) -> D::Pattern

Return the pattern of the dimension

pub fn ndim(&self) -> usize

Returns rank (ndim) of the dimension

pub fn raw_dim(&self) -> D

Returns the raw dimension D

pub fn shape(&self) -> &[usize]

Returns a reference to the current dimension, as a slice.

pub fn into_owned(self) -> QkvBase<OwnedRepr<A>, D>

where A: Clone, S: Data,

pub fn to_owned(&self) -> QkvBase<OwnedRepr<A>, D>

where A: Clone, S: Data,

pub fn into_shared(self) -> QkvBase<OwnedArcRepr<A>, D>

where A: Clone, S: DataOwned,

pub fn to_shared(&self) -> QkvBase<OwnedArcRepr<A>, D>

where A: Clone, S: DataShared,

pub fn view(&self) -> QkvBase<ViewRepr<&A>, D>

where S: Data,

pub fn view_mut(&mut self) -> QkvBase<ViewRepr<&mut A>, D>

where S: DataMut,

impl<A, S, D> QkvBase<S, D>

where D: Dimension, S: RawData<Elem = A>, A: Clone,

pub fn attention( &self, dropout: Option<f64>, mask: Option<&Array<bool, D>>, ) -> Score<A, D>

where A: ComplexFloat + ScalarOperand, S: Data, ArrayBase<S, D>: for<'a> Dot<ArrayView<'a, A, D>, Output = Array<A, D>>, Array<A, D>: Dot<ArrayBase<S, D>, Output = Array<A, D>>,

Computes the Score using scaled dot-product attention.

Trait Implementations

impl<A, S, D> Borrow<QkvBase<S, D>> for AttentionHead<A, D, S>

where D: Dimension, S: RawData<Elem = A>,

fn borrow(&self) -> &QkvBase<S, D>

Immutably borrows from an owned value.

impl<A, S, D> BorrowMut<QkvBase<S, D>> for AttentionHead<A, D, S>

where D: Dimension, S: RawData<Elem = A>,

fn borrow_mut(&mut self) -> &mut QkvBase<S, D>

Mutably borrows from an owned value.

impl<A, S, D> Clone for QkvBase<S, D>

where D: Dimension, S: RawDataClone<Elem = A>,

fn clone(&self) -> Self

Returns a copy of the value.

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

Performs copy-assignment from source.

impl<A, S, D> Default for QkvBase<S, D>

where A: Default, D: Dimension, S: DataOwned<Elem = A>,

fn default() -> Self

Returns the “default value” for a type.

impl<A, S, D> From<QkvBase<S, D>> for AttentionHead<A, D, S>

where D: Dimension, S: RawData<Elem = A>,

fn from(params: QkvBase<S, D>) -> Self

Converts to this type from the input type.

impl<A, S, T, D, E, F> Matmul<ArrayBase<T, E>> for QkvBase<S, D>

where A: LinalgScalar, D: Dimension, E: Dimension, F: Dimension, S: Data<Elem = A>, T: Data<Elem = A>, ArrayBase<S, D>: Dot<ArrayBase<T, E>, Output = Array<A, F>>,

type Output = QkvBase<OwnedRepr<A>, F>

fn matmul(&self, rhs: &ArrayBase<T, E>) -> Self::Output

impl<A, S, T, D, E, F> Matmul<QkvBase<T, E>> for QkvBase<S, D>

where A: LinalgScalar, D: Dimension, E: Dimension, F: Dimension, S: Data<Elem = A>, T: Data<Elem = A>, ArrayBase<S, D>: Dot<ArrayBase<T, E>, Output = Array<A, F>>,

type Output = QkvBase<OwnedRepr<A>, F>

fn matmul(&self, rhs: &QkvBase<T, E>) -> Self::Output

impl<A, B, S, D, S2, D2> PartialEq<(ArrayBase<S2, D2>, ArrayBase<S2, D2>, ArrayBase<S2, D2>)> for QkvBase<S, D>

where A: PartialEq, B: PartialEq, D: Dimension, S: Data<Elem = A>, S2: Data<Elem = B>, D2: Dimension, ArrayBase<S, D>: PartialEq<ArrayBase<S2, D2>>,

fn eq(&self, (q, k, v): &(ArrayBase<S2, D2>, A, B)) -> 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<A, B, S, D, S2, D2> PartialEq<ArrayBase<S2, D2>> for QkvBase<S, D>

where A: PartialEq, B: PartialEq, D: Dimension, S: Data<Elem = A>, S2: Data<Elem = B>, D2: Dimension, ArrayBase<S, D>: PartialEq<ArrayBase<S2, D2>>,

fn eq(&self, other: &ArrayBase<S2, D2>) -> 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<A, S, D> PartialEq for QkvBase<S, D>

where A: PartialEq, D: Dimension, S: Data<Elem = A>,

fn eq(&self, other: &Self) -> 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<A, S, D> Copy for QkvBase<S, D>

where D: Copy + Dimension, S: Copy + RawDataClone<Elem = A>,