cellflow.networks.ConditionalVelocityField

class cellflow.networks.ConditionalVelocityField(output_dim, max_combination_length, condition_mode='deterministic', regularization=1.0, condition_embedding_dim=32, covariates_not_pooled=<factory>, pooling='attention_token', pooling_kwargs=<factory>, layers_before_pool=<factory>, layers_after_pool=<factory>, cond_output_dropout=0.0, condition_dropout_prob=0.0, condition_null='zero_embedding', mask_value=0.0, condition_encoder_kwargs=<factory>, act_fn=<PjitFunction of <function silu>>, time_freqs=1024, time_max_period=10000, time_encoder_dims=(1024, 1024, 1024), time_encoder_dropout=0.0, hidden_dims=(1024, 1024, 1024), hidden_dropout=0.0, conditioning='concatenation', conditioning_kwargs=<factory>, decoder_dims=(1024, 1024, 1024), decoder_dropout=0.0, layer_norm_before_concatenation=False, linear_projection_before_concatenation=False, parent=<flax.linen.module._Sentinel object>, name=None)

Parameterized neural vector field with conditions.

Parameters:

• output_dim (int) – Dimensionality of the output.

• max_combination_length (int) – Maximum number of covariates in a combination.

• 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.

• condition_embedding_dim (int) – Dimensions of the condition embedding.

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

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

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

• layers_before_pool (dict[str, Sequence[dict[str, Any]]] | 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.

• cond_output_dropout (float) – Dropout rate for the last layer of the condition encoder.

• condition_dropout_prob (float) – Probability of dropping the whole condition during training so that the network also learns an unconditional velocity field. This enables classifier-free guidance at inference time via the force_uncond argument of __call__(). A value of 0.0 (the default) disables condition dropout and reproduces the standard conditional behavior.

• condition_null (Literal['zero_embedding', 'mask_value']) –

  How the condition is nulled when it is dropped (both during training via condition_dropout_prob and at inference via force_uncond):

  • 'zero_embedding': zero the condition embedding after the encoder (the default).

  • 'mask_value': fill the raw condition inputs with mask_value and route the resulting fully-masked set through the encoder.

• condition_encoder_kwargs (dict[str, Any]) – Keyword arguments for the condition encoder.

• act_fn (Callable[[Array], Array]) – Activation function.

• time_freqs (int) – Frequency of the cyclical time encoding.

• time_max_period (int) – Controls the minimum frequency of the time embeddings.

• time_encoder_dims (Sequence[int]) – Dimensions of the time embedding.

• time_encoder_dropout (float) – Dropout rate for the time embedding.

• hidden_dims (Sequence[int]) – Dimensions of the hidden layers.

• hidden_dropout (float) – Dropout rate for the hidden layers.

• conditioning (Literal['concatenation', 'film', 'resnet']) –

  Conditioning method, should be one of:

  • 'concatenation': Concatenate the time, data, and condition embeddings.

  • 'film': Use FiLM conditioning, i.e. learn FiLM weights from time and condition embedding to scale the data embeddings.

  • 'resnet': Use residual conditioning.

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

• decoder_dims (Sequence[int]) – Dimensions of the output layers.

• decoder_dropout (float) – Dropout rate for the output layers.

• layer_norm_before_concatenation (bool) – If True, applies layer normalization before concatenating the embedded time, embedded data, and condition embeddings.

• linear_projection_before_concatenation (bool) – If True, applies a linear projection before concatenating the embedded time, embedded data.

• mask_value (float)

• parent (Module | Scope | _Sentinel | None)

• name (str | None)

Returns:

Output of the neural vector field.

Methods

act_fn()	SiLU (aka swish) activation function.
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 the training state.
get_condition_embedding(condition)	Get the embedding of the condition.
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 network.
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__(t, x_t, cond, encoder_noise[, ...])	Call self as a function.

Attributes

cond_output_dropout
condition_dropout_prob
condition_embedding_dim
condition_mode
condition_null
conditioning
decoder	The decoder used.
decoder_dims
decoder_dropout
hidden_dims
hidden_dropout
layer_norm_before_concatenation
linear_projection_before_concatenation
mask_value
name
output_dims	Dimensions of the output layers.
parent
path	Get the path of this Module.
pooling
regularization
scope
time_encoder	The time encoder used.
time_encoder_dims
time_encoder_dropout
time_freqs
time_max_period
variables	Returns the variables in this module.
x_encoder	The x encoder used.
output_dim
max_combination_length
covariates_not_pooled
pooling_kwargs
layers_before_pool
layers_after_pool
condition_encoder_kwargs
conditioning_kwargs