AQIT 0.1.0
Loading...
Searching...
No Matches
aquin.compute.sae.SparseAutoencoder Class Reference
Inheritance diagram for aquin.compute.sae.SparseAutoencoder:
Collaboration diagram for aquin.compute.sae.SparseAutoencoder:

Public Member Functions

 __init__ (self, int d_model=2048, int n_features=16384)
torch.Tensor encode (self, torch.Tensor x)
torch.Tensor decode (self, torch.Tensor f)
 forward (self, torch.Tensor x)
list[tuple[int, float]] get_top_features (self, torch.Tensor x, int k=10)
None save (self, str|Path path)
"SparseAutoencoder" load (cls, str|Path path, str device="cuda")

Public Attributes

 d_model = d_model
 n_features = n_features
 b_pre = nn.Parameter(torch.zeros(d_model))
 W_enc = nn.Parameter(torch.nn.init.kaiming_uniform_(torch.empty(d_model, n_features)))
 b_enc = nn.Parameter(torch.zeros(n_features))
 W_dec = nn.Parameter(torch.nn.init.kaiming_uniform_(torch.empty(n_features, d_model)))
 b_dec = nn.Parameter(torch.zeros(d_model))

Protected Member Functions

None _normalise_decoder (self)
torch.Tensor|None _find_tensor (cls, dict tensors, *str names)
tuple[int, int] _infer_dims (cls, dict tensors)
tuple[dict[str, Any], dict[str, torch.Tensor]] _coerce_checkpoint (cls, Any ckpt)
int|None _meta_int (cls, dict[str, Any] meta, *str keys)
dict[str, torch.Tensor] _remap_tensors (cls, dict[str, torch.Tensor] tensors, "SparseAutoencoder" sae)
"SparseAutoencoder" _load_safetensors (cls, Path path, str device="cuda")

Detailed Description

Definition at line 18 of file sae.py.

Constructor & Destructor Documentation

◆ __init__()

aquin.compute.sae.SparseAutoencoder.__init__ ( self,
int d_model = 2048,
int n_features = 16384 )

Definition at line 19 of file sae.py.

References __init__().

Referenced by __init__().

Member Function Documentation

◆ _coerce_checkpoint()

tuple[dict[str, Any], dict[str, torch.Tensor]] aquin.compute.sae.SparseAutoencoder._coerce_checkpoint ( cls,
Any ckpt )
protected

Definition at line 89 of file sae.py.

Referenced by load().

◆ _find_tensor()

torch.Tensor | None aquin.compute.sae.SparseAutoencoder._find_tensor ( cls,
dict tensors,
*str names )
protected

Definition at line 57 of file sae.py.

Referenced by _infer_dims(), _remap_tensors(), and load().

◆ _infer_dims()

tuple[int, int] aquin.compute.sae.SparseAutoencoder._infer_dims ( cls,
dict tensors )
protected

Definition at line 72 of file sae.py.

References _find_tensor().

Referenced by load().

◆ _load_safetensors()

"SparseAutoencoder" aquin.compute.sae.SparseAutoencoder._load_safetensors ( cls,
Path path,
str device = "cuda" )
protected

Definition at line 224 of file sae.py.

Referenced by load().

◆ _meta_int()

int | None aquin.compute.sae.SparseAutoencoder._meta_int ( cls,
dict[str, Any] meta,
*str keys )
protected

Definition at line 116 of file sae.py.

Referenced by load().

◆ _normalise_decoder()

None aquin.compute.sae.SparseAutoencoder._normalise_decoder ( self)
protected

Definition at line 30 of file sae.py.

References W_dec.

◆ _remap_tensors()

dict[str, torch.Tensor] aquin.compute.sae.SparseAutoencoder._remap_tensors ( cls,
dict[str, torch.Tensor] tensors,
"SparseAutoencoder" sae )
protected

Definition at line 123 of file sae.py.

References _find_tensor().

Referenced by load().

◆ decode()

torch.Tensor aquin.compute.sae.SparseAutoencoder.decode ( self,
torch.Tensor f )

Definition at line 38 of file sae.py.

References b_dec, and W_dec.

Referenced by forward().

◆ encode()

torch.Tensor aquin.compute.sae.SparseAutoencoder.encode ( self,
torch.Tensor x )

Definition at line 35 of file sae.py.

References b_enc, b_pre, and W_enc.

Referenced by forward(), and get_top_features().

◆ forward()

aquin.compute.sae.SparseAutoencoder.forward ( self,
torch.Tensor x )

Definition at line 41 of file sae.py.

References decode(), and encode().

◆ get_top_features()

list[tuple[int, float]] aquin.compute.sae.SparseAutoencoder.get_top_features ( self,
torch.Tensor x,
int k = 10 )

Definition at line 46 of file sae.py.

References encode().

◆ load()

"SparseAutoencoder" aquin.compute.sae.SparseAutoencoder.load ( cls,
str | Path path,
str device = "cuda" )

Definition at line 160 of file sae.py.

References _coerce_checkpoint(), _find_tensor(), _infer_dims(), _load_safetensors(), _meta_int(), and _remap_tensors().

◆ save()

None aquin.compute.sae.SparseAutoencoder.save ( self,
str | Path path )

Definition at line 51 of file sae.py.

References d_model, and n_features.

Member Data Documentation

◆ b_dec

aquin.compute.sae.SparseAutoencoder.b_dec = nn.Parameter(torch.zeros(d_model))

Definition at line 27 of file sae.py.

Referenced by decode().

◆ b_enc

aquin.compute.sae.SparseAutoencoder.b_enc = nn.Parameter(torch.zeros(n_features))

Definition at line 25 of file sae.py.

Referenced by encode().

◆ b_pre

aquin.compute.sae.SparseAutoencoder.b_pre = nn.Parameter(torch.zeros(d_model))

Definition at line 23 of file sae.py.

Referenced by encode().

◆ d_model

aquin.compute.sae.SparseAutoencoder.d_model = d_model

Definition at line 21 of file sae.py.

Referenced by save().

◆ n_features

aquin.compute.sae.SparseAutoencoder.n_features = n_features

Definition at line 22 of file sae.py.

Referenced by save().

◆ W_dec

aquin.compute.sae.SparseAutoencoder.W_dec = nn.Parameter(torch.nn.init.kaiming_uniform_(torch.empty(n_features, d_model)))

Definition at line 26 of file sae.py.

Referenced by _normalise_decoder(), and decode().

◆ W_enc

aquin.compute.sae.SparseAutoencoder.W_enc = nn.Parameter(torch.nn.init.kaiming_uniform_(torch.empty(d_model, n_features)))

Definition at line 24 of file sae.py.

Referenced by encode().

The documentation for this class was generated from the following file: