"""Utilities functions for the layer submodules."""

import torch

from . import CoordConv


def get_conv_layer(conv_type: str = "standard") -> torch.nn.Module:
    """Return a conv layer based on the passed in string name."""
    if conv_type == "standard":
        conv_layer = torch.nn.Conv2d
    elif conv_type == "coord":
        conv_layer = CoordConv
    elif conv_type == "3d":
        conv_layer = torch.nn.Conv3d
    else:
        raise ValueError(f"{conv_type} is not a recognized Conv method")
    return conv_layer