dreams.models.layers package

Submodules

dreams.models.layers.feed_forward module

class dreams.models.layers.feed_forward.FeedForward(in_dim, out_dim, hidden_dim, depth=None, act_last=True, act=<class 'torch.nn.modules.activation.ReLU'>, bias=True, dropout=0)

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

dreams.models.layers.fourier_features module

class dreams.models.layers.fourier_features.FourierFeatures(strategy, x_min, x_max, trainable=True, funcs='both', sigma=10, num_freqs=512)

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

num_features()

dreams.models.layers.fp_growing module

Copypasted from https://github.com/samgoldman97/mist/blob/main_v2/src/mist/models/modules.py

class dreams.models.layers.fp_growing.FPGrowingModule(hidden_input_dim: int = 256, final_target_dim: int = 4096, num_splits=4, reduce_factor=2)

Bases: Module

FPGrowingModule.

Accept an input hidden dim and progressively grow by powers of 2 s.t.

We eventually get to the final output size…

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(hidden)

forward.

Return dict mapping output dim to the

dreams.models.layers.tnq_transformer module

Module contents