Struct MemoryAttentionCompiled 

pub struct MemoryAttentionCompiled {
    pub n_img: usize,
    pub max_n_mem: usize,
    pub d_model: usize,
    pub kv_in_dim: usize,
    pub max_obj_ptr_tokens: usize,
    /* private fields */
}

Fields

n_img: usize

max_n_mem: usize

d_model: usize

kv_in_dim: usize

max_obj_ptr_tokens: usize

Implementations

impl MemoryAttentionCompiled

pub fn compile( w: &Sam2MemoryAttentionWeights, n_img: usize, max_n_mem: usize, max_obj_ptr_tokens: usize, device: Device, ) -> Result<MemoryAttentionCompiled, Error>

pub fn compile_with_profile( w: &Sam2MemoryAttentionWeights, n_img: usize, max_n_mem: usize, max_obj_ptr_tokens: usize, device: Device, profile: &CompileProfile, ) -> Result<MemoryAttentionCompiled, Error>

pub fn compile_in_graph_rope( w: &Sam2MemoryAttentionWeights, n_img: usize, max_n_mem: usize, max_obj_ptr_tokens: usize, device: Device, ) -> Result<MemoryAttentionCompiled, Error>

One compiled graph per layer (Op::AxialRope2d in-graph). Fusion is disabled so FuseAttentionBlock cannot reorder RoPE relative to attention.

pub fn compile_in_graph_rope_with_profile( w: &Sam2MemoryAttentionWeights, n_img: usize, max_n_mem: usize, max_obj_ptr_tokens: usize, device: Device, profile: &CompileProfile, ) -> Result<MemoryAttentionCompiled, Error>

pub fn run( &mut self, curr: &[f32], curr_pos: &[f32], memory: &[f32], memory_pos: &[f32], active_n_mem: usize, num_obj_ptr_tokens: usize, ) -> Result<Vec<f32>, Error>