Stacked MLP Model

A set of classes defining a simpler (stacked) MLP models

class mentat_lss.models.stacked_mlp.single_mlp(config_dict, is_cross_spectra: bool)

Class defining a single independent mlp network

forward(input_params)

Passes an input tensor through the network

class mentat_lss.models.stacked_mlp.stacked_mlp(config_dict)

Class defining a stack of single_mlp objects, one for each portion of the power spectrum output

forward(input_params, net_idx=None)

Passes an input tensor through the network

organize_parameters(input_params)

Organizes input cosmology + bias parameters into a form the rest of the network expects

Parameters:

input_params – tensor of input parameters with shape [batch, num_cosmo_params*(num_nuisance_params*num_zbins*num_tracers)]

Returns:

tensor of input parameters with shape [batch, num_spectra*num_zbins, num_cosmo_params + (2*self.num_nuisance_params)].

The bias parameters are split corresponding to their respective redshift / tracer bin

Return type:

organized_params