nn

This module contains lucidlearn’s make_* functions for assembling deep-learning models. These range from simple building blocks like make_linear to specialized architectures like make_spectrum_encoder, which embeds peaks from tandem mass spectrometry (MS/MS) experiments so that, given the right model head, molecular properties or structures can be predicted from the spectra.

source

make_property_predictor

def make_property_predictor(
    key, d:int=512, ldmin:float=0.0031622776601683794, ldmax:float=1995.2623149688789, hidden_layer:int=512,
    ffdim:int=512, nhead:int=32, nlayers:int=6, dout:int=10, act_fn:jax._src.custom_derivatives.custom_jvp=relu,
    initializer:function=he_normal, attention_backend:NoneType=None, block_size:int=128, bias:bool=True,
    mask_fn:function=mask_fn
):

source

make_spectrum_encoder

def make_spectrum_encoder(
    key, d:int=512, ldmin:float=0.0031622776601683794, ldmax:float=1995.2623149688789, hidden_layer:int=512,
    ffdim:int=512, nhead:int=32, nlayers:int=6, act_fn:jax._src.custom_derivatives.custom_jvp=relu,
    initializer:function=he_normal, attention_backend:NoneType=None, block_size:int=128, mask_fn:function=mask_fn
):

source

make_transformer_block

def make_transformer_block(
    key, nembd, nhead, ffdim:NoneType=None, attention_backend:NoneType=None, block_size:int=128,
    act_fn:jax._src.custom_derivatives.custom_jvp=relu, initializer:function=he_normal, bias:bool=True
):

source

make_multi_head_attention

def make_multi_head_attention(
    key, nembd, nhead, attention_backend:NoneType=None, block_size:int=128, initializer:function=he_normal
):

source

blockwise_attention

def blockwise_attention(
    q, k, v, mask:NoneType=None, block_size:int=128
):

q: (B, N, T, H), k: (B, N, T, H), v: (B, N, T, H)

source

mask_fn

def mask_fn(
    x
):

source

make_peak_embedding

def make_peak_embedding(
    key, d:int=512, hidden_layer:int=512, ldmin:float=0.0031622776601683794, ldmax:float=1995.2623149688789,
    act_fn:jax._src.custom_derivatives.custom_jvp=relu, initializer:function=he_normal, bias:bool=True
):

source

make_sinusoidal_embedding

def make_sinusoidal_embedding(
    d:int=512, ldmin:float=0.0031622776601683794, ldmax:float=1995.2623149688789
):

source

make_layernorm2d

def make_layernorm2d(
    key, num_features, eps:float=1e-05, offset:bool=True, ndims:int=4
):

source

make_layernorm1d

def make_layernorm1d(
    key, num_features, eps:float=1e-05, offset:bool=True
):

source

make_squeeze

def make_squeeze(
    key, dims:NoneType=None
):

source

make_conv2d

def make_conv2d(
    key, cho, chi, ks:tuple=(3, 3), padding:str='SAME', strides:tuple=(1, 1),
    act_fn:jax._src.custom_derivatives.custom_jvp=relu, initializer:function=he_normal, bias:bool=True,
    act:bool=True
):

source

make_ffn

def make_ffn(
    key, din, dout, hidden_layer:int=512, act_fn:jax._src.custom_derivatives.custom_jvp=relu,
    initializer:function=he_normal, bias:bool=True
):

source

make_linear

def make_linear(
    key, fan_in, fan_out, act_fn:jax._src.custom_derivatives.custom_jvp=relu, initializer:function=he_normal,
    bias:bool=True, act:bool=True
):