cellflow.networks.TokenAttentionPooling

class cellflow.networks.TokenAttentionPooling(num_heads=8, qkv_dim=64, dropout_rate=0.0, num_layers=1, transformer_block=False, layer_norm=False, ff_dim=None, act_fn=<PjitFunction of <function silu>>, parent=<flax.linen.module._Sentinel object>, name=None)

Multi-head attention which aggregates sets by learning a token.

A learnable [CLS] token is prepended to the set and num_layers self-attention layers are applied over {CLS, *tokens} (BERT/ViT-style); the CLS token’s final state is returned as the aggregate. num_layers=1 with transformer_block=False reproduces the original single-attention pooling. Set num_layers>=2 (+ transformer_block=True) to make the condition encoder a transformer where every token attends to every other before the CLS read-out.

Parameters:

• num_heads (int) – Number of attention heads (per layer).

• qkv_dim (int) – Dimensionality of the query, key, and value (per layer).

• dropout_rate (float) – Dropout rate.

• num_layers (int) – Number of stacked self-attention layers over {CLS, *tokens}.

• transformer_block (bool) – Whether each layer is a full transformer block (adds a feed-forward sublayer with residual).

• layer_norm (bool) – Whether to use layer normalization inside each transformer block.

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

• ff_dim (int | None)

• parent (Module | Scope | _Sentinel | None)

• name (str | None)

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.
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()	Initializes a Module lazily (similar to a lazy __init__).
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__(x[, mask, training])	Forward pass.

Attributes

dropout_rate
ff_dim
layer_norm
name
num_heads
num_layers
parent
path	Get the path of this Module.
qkv_dim
scope
takes_attention_mask	Whether this block's __call__ accepts an attention mask as its second positional argument.
transformer_block
variables	Returns the variables in this module.