kaira.models.components.AFModule

Inheritance diagram of AFModule

Inheritance diagram for AFModule

class kaira.models.components.AFModule(N, csi_length, *args: Any, **kwargs: Any)[source]

Bases: ChannelAwareBaseModel

AFModule: Attention-Feature Module [Xu et al., 2021].

This module implements a an attention mechanism that recalibrates feature maps by explicitly modeling interdependencies between channel state information and the input features. This module allows the same model to be used during training and testing across channels with different signal-to-noise ratio without significant performance degradation.

Methods

__init__

Initialize the AFModule.

create_csi_for_submodules

Create appropriate CSI tensors for multiple submodules.

extract_csi_features

Extract common features from CSI tensor for analysis.

extract_csi_from_channel_output

Extract CSI from channel output if available.

format_csi_for_channel

Format CSI tensor for passing to channels that expect specific formats.

forward

Forward pass through the AFModule.

forward_csi_to_submodules

Helper method to consistently pass CSI to submodules.

normalize_csi

Normalize CSI values to a specified range.

transform_csi

Transform CSI tensor to match target shape requirements.

validate_csi

Validate and ensure CSI tensor is in the correct format.
__init__(N, csi_length, *args: Any, **kwargs: Any)[source]

Initialize the AFModule.

Parameters:

  • N (int) – The number of input and output features.

  • csi_length (int) – The length of the channel state information.

  • *args – Variable positional arguments passed to the base class.

  • **kwargs – Variable keyword arguments passed to the base class.

forward(x: Tensor, csi: Tensor, *args: Any, **kwargs: Any) → Tensor[source]

Forward pass through the AFModule.

Parameters:

  • x (torch.Tensor) – The input tensor.

  • csi (torch.Tensor) – Channel State Information tensor.

  • *args – Additional positional arguments (unused).

  • **kwargs – Additional keyword arguments (unused).

Returns:

The output tensor after applying the attention mechanism.

Return type:

torch.Tensor

create_csi_for_submodules(csi: Tensor, num_modules: int) → List[Tensor]

Create appropriate CSI tensors for multiple submodules.

Parameters:

  • csi (torch.Tensor) – Original CSI tensor

  • num_modules (int) – Number of submodules that need CSI

Returns:

List of CSI tensors for each submodule

Return type:

List[torch.Tensor]

extract_csi_features(csi: Tensor) → Dict[str, Tensor]

Extract common features from CSI tensor for analysis.

Parameters:

csi (torch.Tensor) – The CSI tensor to analyze

Returns:

Dictionary containing extracted features

Return type:

Dict[str, torch.Tensor]

static extract_csi_from_channel_output(channel_output: Any) → Tensor | None

Extract CSI from channel output if available.

Some channels return both the transmitted signal and CSI information. This static method provides a standardized way to extract CSI from various channel output formats.

Parameters:

channel_output – Output from a channel, which may contain CSI

Returns:

Extracted CSI tensor if available, None otherwise

Return type:

Optional[torch.Tensor]

static format_csi_for_channel(csi: Tensor, channel_format: str = 'tensor') → Any

Format CSI tensor for passing to channels that expect specific formats.

Parameters:

  • csi (torch.Tensor) – CSI tensor to format

  • channel_format (str) – Expected format (“tensor”, “dict”, “kwargs”)

Returns:

Formatted CSI in the requested format

Return type:

Any

forward_csi_to_submodules(csi: Tensor, modules: List[BaseModel], *args, **kwargs) → List[Any]

Helper method to consistently pass CSI to submodules.

This method facilitates passing CSI to multiple submodules that require channel state information, ensuring consistent handling across the model.

Parameters:

  • csi (torch.Tensor) – Channel state information tensor

  • modules (List[BaseModel]) – List of modules to apply

  • *args – Positional arguments to pass to modules

  • **kwargs – Keyword arguments to pass to modules

Returns:

List of outputs from each module

Return type:

List[Any]

normalize_csi(csi: Tensor, method: str = 'minmax', target_range: tuple = (0.0, 1.0)) → Tensor

Normalize CSI values to a specified range.

Parameters:

  • csi (torch.Tensor) – The CSI tensor to normalize

  • method (str) – Normalization method. Options: “minmax”, “zscore”, “none”

  • target_range (tuple) – Target range for minmax normalization (min, max)

Returns:

Normalized CSI tensor

Return type:

torch.Tensor

Raises:

ValueError – If normalization method is not supported

transform_csi(csi: Tensor, target_shape: Size) → Tensor

Transform CSI tensor to match target shape requirements.

Parameters:

  • csi (torch.Tensor) – The CSI tensor to transform

  • target_shape (torch.Size) – Target shape for the CSI tensor

Returns:

Transformed CSI tensor

Return type:

torch.Tensor

validate_csi(csi: Tensor, expected_shape: Size | None = None) → Tensor

Validate and ensure CSI tensor is in the correct format.

Parameters:

  • csi (torch.Tensor) – The CSI tensor to validate

  • expected_shape (Optional[torch.Size]) – Expected shape for the CSI tensor. If None, uses cached shape or infers from tensor.

Returns:

Validated CSI tensor

Return type:

torch.Tensor

Raises:

  • ValueError – If CSI tensor is invalid or has incorrect shape

  • TypeError – If CSI is not a tensor