xlstm_jax.models.xlstm_clean.components.init#

Functions#

`bias_linspace_init`(start, end)	Linearly spaced bias init across dimensions.
`small_init`(dim)	Return initializer that creates a tensor with values according to the method described in:
`wang_init`(dim, num_blocks)	Adopted from EleutherAI/gpt-neox.
`uniform_init`(min_val, max_val)	Uniform initializer.

Module Contents#

xlstm_jax.models.xlstm_clean.components.init.bias_linspace_init(start, end)#

Linearly spaced bias init across dimensions.

Parameters:

start (float)
end (float)

Return type:

callable

xlstm_jax.models.xlstm_clean.components.init.small_init(dim)#

Return initializer that creates a tensor with values according to the method described in: “Transformers without Tears: Improving the Normalization of Self-Attention”, Nguyen, T. & Salazar, J. (2019).

Adopted from EleutherAI/gpt-neox.

Parameters:: dim (int)
Return type:: callable

xlstm_jax.models.xlstm_clean.components.init.wang_init(dim, num_blocks)#

Adopted from EleutherAI/gpt-neox.

Parameters:

dim (int)
num_blocks (int)

Return type:

callable

xlstm_jax.models.xlstm_clean.components.init.uniform_init(min_val, max_val)#

Uniform initializer.

Parameters:

min_val (float)
max_val (float)

Return type:

callable