Skip to content

lens

To see all the implemented Lens, simply import Lens and print Lens.registry.

flowchart LR
    Lens
    subgraph Registry
       LensA
    end

    Lens-->LensA

Lens 

Lens(unembed, bias, n_head, d_model, d_vocab)

Bases: Module

It is important to remember that subclasses of the Lens cannot take separate arguments as they will be ignored at runtime for instantiation in the TransformerAttnInterface.

Parameters:

  • unembed –
  • bias –
  • n_head –
  • d_model –
  • d_vocab –

get_lens classmethod 

get_lens(name: str) -> type[Lens]

This takes the name of the lens and queries the registry to grab the corresponding Lens subclas.

Parameters:

  • name (str) –

    The name of the child Lens implementation.

Returns:

  • type[Lens] –

    Subclass Lens implementation.

LensA 

LensA(unembed, bias, n_head, d_model, d_vocab)

Bases: Lens

get_lens classmethod 

get_lens(name: str) -> type[Lens]

This takes the name of the lens and queries the registry to grab the corresponding Lens subclas.

Parameters:

  • name (str) –

    The name of the child Lens implementation.

Returns:

  • type[Lens] –

    Subclass Lens implementation.

forward 

forward(input_tensor)

Performs a forward pass through the LensA model.

Parameters:

  • input_tensor (Tensor) –

    Input tensor of shape (batch_size, pos, n_head, d_model).

Returns:

  • –

    torch.tensor: Output tensor of shape (batch_size, pos, d_vocab) after processing through

  • –

    the linear layers and summing across the attention heads.