cellflow.networks.SelfAttention.__call__

SelfAttention.__call__(x, mask=None, training=True)

Forward pass.

Parameters:

• x (Array) – Input tensor of shape (batch_size, set_size, input_dim) or (batch_size, input_dim).

• mask (Array | None) – Mask tensor of shape (batch_size, 1 | num_heads, set_size, set_size).

• training (bool) – Whether the model is in training mode.

Returns:

: Output tensor of shape (batch_size, set_size, input_dim).