Struct VitEncoder 

pub struct VitEncoder<B: Backend> { /* private fields */ }

Vision Transformer encoder.

Maps images to patch-level representations via:

1. Patch embedding (linear projection of flattened patches)
2. 2D Rotary Position Encoding
3. Stack of transformer blocks (self-attention + MLP)
4. Final layer normalization

Output shape: [batch, num_patches, embed_dim]

Implementations

impl<B: Backend> VitEncoder<B>

pub fn forward(&self, images: &Tensor<B, 4>) -> Representation<B>

Forward pass: image → representation.

Arguments

• images - Input images. Shape: [batch, channels, height, width]

Returns

Patch-level representations. Shape: [batch, num_patches, embed_dim]

pub fn forward_visible_tokens( &self, images: &Tensor<B, 4>, visible_indices: &[usize], ) -> Representation<B>

Encode only the visible patch tokens for strict JEPA context encoding.

The image is patchified and position-encoded using the full grid so the surviving tokens retain their real flattened positions, then masked tokens are removed before self-attention runs.

pub fn load_named_tensors( self, tensors: &HashMap<String, TensorData>, ) -> Result<Self, VitLoadError>

Load a ViT encoder from a map of burn-style parameter names to tensor data.

Expected parameter names match the burn module record layout, for example patch_embed.projection.weight and blocks.0.attn.out_proj.bias.

pub fn ema_update_from(self, online: &Self, ema: &Ema, step: usize) -> Self

Update this encoder toward an online encoder using EMA.

The returned encoder preserves the gradient setting of the target encoder parameters while detaching the blended tensors from any active autodiff graph.

Trait Implementations

impl<B> AutodiffModule<B> for VitEncoder<B>

where B: AutodiffBackend + Backend, <B as AutodiffBackend>::InnerBackend: Backend,

type InnerModule = VitEncoder<<B as AutodiffBackend>::InnerBackend>

Inner module without auto-differentiation.

fn valid(&self) -> Self::InnerModule

Get the same module, but on the inner backend without auto-differentiation.

impl<B: Backend> Clone for VitEncoder<B>

fn clone(&self) -> Self

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl<B: Debug + Backend> Debug for VitEncoder<B>

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

Formats the value using the given formatter.

impl<B: Backend> Display for VitEncoder<B>

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

Formats the value using the given formatter.

impl<B: Backend> Encoder<B> for VitEncoder<B>

type Input = Tensor<B, 4>

The type of input this encoder accepts.

fn encode(&self, input: &Self::Input) -> Representation<B>

Encode input into a representation.

fn embed_dim(&self) -> usize

Get the output embedding dimension.

impl<B: Backend> Module<B> for VitEncoder<B>

type Record = VitEncoderRecord<B>

Type to save and load the module.

fn load_record(self, record: Self::Record) -> Self

Load the module state from a record.

fn into_record(self) -> Self::Record

Convert the module into a record containing the state.

fn num_params(&self) -> usize

Get the number of parameters the module has, including all of its sub-modules.

fn visit<Visitor: ModuleVisitor<B>>(&self, visitor: &mut Visitor)

Visit each tensor parameter in the module with a visitor.

fn map<Mapper: ModuleMapper<B>>(self, mapper: &mut Mapper) -> Self

Map each tensor parameter in the module with a mapper.

fn collect_devices(&self, devices: Devices<B>) -> Devices<B>

Return all the devices found in the underneath module tree added to the given vector without duplicates.

fn to_device(self, device: &B::Device) -> Self

Move the module and all of its sub-modules to the given device.

fn fork(self, device: &B::Device) -> Self

Fork the module and all of its sub-modules to the given device.

fn devices(&self) -> Vec<<B as Backend>::Device>

Return all the devices found in the underneath module tree without duplicates.

fn no_grad(self) -> Self

Each tensor in the module tree will not require grad.

fn save_file<FR, PB>( self, file_path: PB, recorder: &FR, ) -> Result<(), RecorderError>

where FR: FileRecorder<B>, PB: Into<PathBuf>,

Save the module to a file using the provided file recorder.

fn load_file<FR, PB>( self, file_path: PB, recorder: &FR, device: &<B as Backend>::Device, ) -> Result<Self, RecorderError>

where FR: FileRecorder<B>, PB: Into<PathBuf>,

Load the module from a file using the provided file recorder.

fn quantize_weights(self, quantizer: &mut Quantizer) -> Self

Quantize the weights of the module.

impl<B: Backend> ModuleDisplay for VitEncoder<B>

fn format(&self, passed_settings: DisplaySettings) -> String

Formats the module with provided display settings.

fn custom_settings(&self) -> Option<DisplaySettings>

Custom display settings for the module.

fn custom_content(&self, _content: Content) -> Option<Content>

Custom attributes for the module.

impl<B: Backend> ModuleDisplayDefault for VitEncoder<B>

fn content(&self, content: Content) -> Option<Content>

Attributes of the module used for display purposes.

fn num_params(&self) -> usize

Gets the number of the parameters of the module.