cellflow.networks.ConditionEncoder

class cellflow.networks.ConditionEncoder(output_dim, condition_mode='deterministic', regularization=0.0, decoder=False, pooling='attention_token', pooling_kwargs=<factory>, covariates_not_pooled=<factory>, layers_before_pool=<factory>, layers_after_pool=<factory>, layers_decoder=<factory>, output_dropout=0.0, mask_value=0.0, parent=<flax.linen.module._Sentinel object>, name=None)[source]

Encoder for conditions represented as sets of perturbations.

Parameters:

  • output_dim (int) – Dimensionality of the output.

  • condition_mode (Literal['deterministic', 'stochastic']) –

    Mode of the encoder, should be one of:

    • 'deterministic': Learns condition encoding point-wise.

    • 'stochastic': Learns a Gaussian distribution for representing conditions.

  • regularization (float) –

    Regularization strength in the latent space:

    • For deterministic mode, it is the strength of the L2 regularization.

    • For stochastic mode, it is the strength of the KL divergence regularization.

  • decoder (bool) – Whether to use a decoder.

  • pooling (Literal['mean', 'attention_token', 'attention_seed']) –

    Pooling method, should be one of:

    • 'mean': Aggregates combinations of covariates by the mean of their learned embeddings.

    • 'attention_token': Aggregates combinations of covariates by an attention mechanism with a token.

    • 'attention_seed': Aggregates combinations of covariates by an attention mechanism with a seed.

  • pooling_kwargs (dict[str, Any]) – Keyword arguments for the pooling method.

  • covariates_not_pooled (Sequence[str]) – Covariates that will escape pooling (should be identical across all set elements).

  • layers_before_pool (Sequence[dict[str, Any]] | dict[str, Sequence[dict[str, Any]]]) – Layers before pooling. Either a sequence of tuples with layer type and parameters or a dictionary with input-specific layers.

  • layers_after_pool (Sequence[dict[str, Any]]) – Layers after pooling.

  • layers_decoder (Sequence[dict[str, Any]]) – Layers for the decoder. Only relevant if 'decoder'=True.

  • mask_value (float) – Value for masked elements used in input conditions.

  • output_dropout (float)

  • parent (Module | Scope | _Sentinel | None)

  • name (str | None)

Methods

apply(variables, *args[, rngs, method, ...])

Applies a module method to variables and returns output and modified variables.

bind(variables, *args[, rngs, mutable])

Creates an interactive Module instance by binding variables and RNGs.

clone(*[, parent, _deep_clone, _reset_names])

Creates a clone of this Module, with optionally updated arguments.

copy(*[, parent, name])

Creates a copy of this Module, with optionally updated arguments.

create_train_state(rng, optimizer, ...)

Create initial training state.

get_variable(col, name[, default])

Retrieves the value of a Variable.

has_rng(name)

Returns true if a PRNGSequence with name name exists.

has_variable(col, name)

Checks if a variable of given collection and name exists in this Module.

init(rngs, *args[, method, mutable, ...])

Initializes a module method with variables and returns modified variables.

init_with_output(rngs, *args[, method, ...])

Initializes a module method with variables and returns output and modified variables.

is_initializing()

Returns True if running under self.init(...) or nn.init(...)().

is_mutable_collection(col)

Returns true if the collection col is mutable.

lazy_init(rngs, *args[, method, mutable])

Initializes a module without computing on an actual input.

make_rng([name])

Returns a new RNG key from a given RNG sequence for this Module.

module_paths(rngs, *args[, show_repeated, ...])

Returns a dictionary mapping module paths to module instances.

param(name, init_fn, *init_args[, unbox])

Declares and returns a parameter in this Module.

perturb(name, value[, collection])

Add an zero-value variable ('perturbation') to the intermediate value.

put_variable(col, name, value)

Updates the value of the given variable if it is mutable, or an error otherwise.

setup()

Initialize the modules.

sow(col, name, value[, reduce_fn, init_fn])

Stores a value in a collection.

tabulate(rngs, *args[, depth, ...])

Creates a summary of the Module represented as a table.

unbind()

Returns an unbound copy of a Module and its variables.

variable(col, name[, init_fn, unbox])

Declares and returns a variable in this Module.

__call__(conditions[, training])

Apply the set encoder.

Attributes

condition_mode

decoder

mask_value

name

output_dropout

parent

path

Get the path of this Module.

pooling

regularization

scope

takes_attention_mask

Whether this block's __call__ accepts an attention mask as its second positional argument.

variables

Returns the variables in this module.

output_dim

pooling_kwargs

covariates_not_pooled

layers_before_pool

layers_after_pool

layers_decoder