elektronn3.modules.axial_attention module

class elektronn3.modules.axial_attention.AxialAttention(*args: Any, **kwargs: Any)

Bases: Module

forward(x)

class elektronn3.modules.axial_attention.AxialImageTransformer(*args: Any, **kwargs: Any)

Bases: Module

forward(x)

class elektronn3.modules.axial_attention.AxialPositionalEmbedding(*args: Any, **kwargs: Any)

Bases: Module

forward(x)

class elektronn3.modules.axial_attention.Deterministic(*args: Any, **kwargs: Any)

Bases: Module

forward(*args, record_rng=False, set_rng=False, **kwargs)

record_rng(*args)

class elektronn3.modules.axial_attention.IrreversibleBlock(*args: Any, **kwargs: Any)

Bases: Module

forward(x, f_args, g_args)

class elektronn3.modules.axial_attention.PermuteToFrom(*args: Any, **kwargs: Any)

Bases: Module

forward(x, **kwargs)

class elektronn3.modules.axial_attention.ReversibleBlock(*args: Any, **kwargs: Any)

Bases: Module

backward_pass(y, dy, f_args={}, g_args={})

forward(x, f_args={}, g_args={})

class elektronn3.modules.axial_attention.ReversibleSequence(*args: Any, **kwargs: Any)

Bases: Module

forward(x, arg_route=(True, True), **kwargs)

class elektronn3.modules.axial_attention.Rezero(*args: Any, **kwargs: Any)

Bases: Module

forward(x)

class elektronn3.modules.axial_attention.SelfAttention(*args: Any, **kwargs: Any)

Bases: Module

forward(x, kv=None)

class elektronn3.modules.axial_attention.Sequential(*args: Any, **kwargs: Any)

Bases: Module

forward(x)

elektronn3.modules.axial_attention.calculate_permutations(num_dimensions, emb_dim)

elektronn3.modules.axial_attention.exists(val)

elektronn3.modules.axial_attention.map_el_ind(arr, ind)

elektronn3.modules.axial_attention.sort_and_return_indices(arr)