diff --git a/src/fastbnns/bnn/inference.py b/src/fastbnns/bnn/inference.py
index c9fce99..61152fe 100644
--- a/src/fastbnns/bnn/inference.py
+++ b/src/fastbnns/bnn/inference.py
@@ -6,13 +6,20 @@
 for that layer if available.
 """
 
+from __future__ import annotations
+
 from collections.abc import Callable, Iterable
-from typing import List, Optional, Union
+from typing import List, Optional, Union, TYPE_CHECKING
 
+import numpy as np
 import torch
 import torch.distributions as dist
 
 
+if TYPE_CHECKING:
+    from fastbnns.bnn.wrappers import BayesianModule
+
+
 class MomentPropagator(torch.nn.Module):
     """Base class for layer propagators."""
 
@@ -200,15 +207,14 @@ def forward(
         return_samples: bool = False,
     ) -> Union[Iterable, tuple]:
         """Propagate moments by averaging over n_samples forward passes of module."""
-        # If the input variance is greater than zero, we'll need to sample the
-        # input as well.
-        if (input[1] > 0.0).all():
+        # If the input variance is not None, we'll need to sample the input as well.
+        if input[1] is None:
+            samples = torch.stack([module(input[0]) for _ in range(self.n_samples)])
+        else:
             input_dist = self.input_sampler(loc=input[0], scale=input[1].sqrt())
             samples = torch.stack(
                 [module(input_dist.sample()) for _ in range(self.n_samples)]
             )
-        else:
-            samples = torch.stack([module(input[0]) for _ in range(self.n_samples)])
 
         if return_samples:
             return type(input)([samples.mean(dim=0), samples.var(dim=0)]), samples
@@ -227,7 +233,7 @@ def __init__(self):
 
     def forward(
         self,
-        module: torch.nn.Module,
+        module: BayesianModule,
         input: Iterable,
     ) -> Iterable:
         """Analytical moment propagation through layer."""
@@ -278,7 +284,7 @@ def __init__(self):
 
     def forward(
         self,
-        module: torch.nn.Module,
+        module: BayesianModule,
         input: Iterable,
     ) -> Iterable:
         """Analytical moment propagation through layer."""
@@ -399,7 +405,7 @@ def __init__(self):
 
     def forward(
         self,
-        module: torch.nn.Module,
+        module: BayesianModule,
         input: Iterable,
         output_size: Optional[List[int]] = None,
     ) -> Iterable:
@@ -470,50 +476,219 @@ class ConvTranspose1d(ConvTransposeNd):
     """Deterministic moment propagation of mean and variance through a ConvTranspose1d layer."""
 
     functional = torch.nn.functional.conv_transpose1d
+    num_spatial_dims = 1
 
     def __init__(self):
         """Initializer for ConvTranspose1d inference module"""
         super().__init__()
-        self.num_spatial_dims = 1
 
 
 class ConvTranspose2d(ConvTransposeNd):
     """Deterministic moment propagation of mean and variance through a ConvTranspose2d layer."""
 
     functional = torch.nn.functional.conv_transpose2d
+    num_spatial_dims = 2
 
     def __init__(self):
         """Initializer for ConvTranspose2d inference module"""
         super().__init__()
-        self.num_spatial_dims = 2
 
 
 class ConvTranspose3d(ConvTransposeNd):
     """Deterministic moment propagation of mean and variance through a ConvTranspose3d layer."""
 
     functional = torch.nn.functional.conv_transpose3d
+    num_spatial_dims = 3
 
     def __init__(self):
         """Initializer for ConvTranspose3d inference module"""
         super().__init__()
-        self.num_spatial_dims = 3
+
+
+class AvgPoolNd(MomentPropagator):
+    """Deterministic moment propagation of mean and variance through AvgPoolNd layers."""
+
+    def __init__(self):
+        """Initializer for AvgPoolNd inference module"""
+        super().__init__()
+
+    def forward(
+        self,
+        module: BayesianModule,
+        input: Iterable,
+    ) -> Iterable:
+        """Analytical moment propagation through layer."""
+        ## Compute analytical result under mean-field approximation following
+        ## https://doi.org/10.48550/arXiv.2402.14532
+        kernel_size = module._module.kernel_size
+        mu = self.functional(
+            input=input[0],
+            kernel_size=kernel_size,
+            stride=module._module.stride,
+            padding=module._module.padding,
+        )
+        n_pool = (
+            kernel_size**self.n_dim
+            if isinstance(kernel_size, int)
+            else torch.prod(torch.tensor(kernel_size))
+        )
+        var = (
+            self.functional(
+                input=input[1],
+                kernel_size=kernel_size,
+                stride=module._module.stride,
+                padding=module._module.padding,
+            )
+            / n_pool
+        )
+
+        return type(input)([mu, var])
+
+
+class AvgPool1d(AvgPoolNd):
+    """Deterministic moment propagation of mean and variance through AvgPool1d layers."""
+
+    functional = torch.nn.functional.avg_pool1d
+    n_dim = 1
+
+    def __init__(self):
+        """Initializer for AvgPool1d inference module"""
+        super().__init__()
+
+
+class AvgPool2d(AvgPoolNd):
+    """Deterministic moment propagation of mean and variance through AvgPool2d layers."""
+
+    functional = torch.nn.functional.avg_pool2d
+    n_dim = 2
+
+    def __init__(self):
+        """Initializer for AvgPool2d inference module"""
+        super().__init__()
+
+
+class AvgPool3d(AvgPoolNd):
+    """Deterministic moment propagation of mean and variance through AvgPool3d layers."""
+
+    functional = torch.nn.functional.avg_pool3d
+    n_dim = 3
+
+    def __init__(self):
+        """Initializer for AvgPool3d inference module"""
+        super().__init__()
+
+
+class ReLUa(MomentPropagator):
+    """Deterministic moment propagation of a normal random variable through a ReLU.
+
+    NOTE: the suffix "a" is added to prevent fastbnns.bnn.wrappers.covert_to_bnn_() from
+    automatically selecting this propagator for ReLU. This is currently desirable
+    since using UnscentedTransform is faster (though less accurate) than analytical propagation.
+    """
+
+    def __init__(self):
+        """Initializer for ReLU inference module"""
+        super().__init__()
+
+    def forward(
+        self,
+        module: BayesianModule,
+        input: Iterable[torch.Tensor],
+    ) -> Iterable[torch.Tensor]:
+        """Analytical moment propagation through layer."""
+        ## Compute analytical result under mean-field approximation following
+        ## https://doi.org/10.48550/arXiv.2402.14532
+        # Compute the mean of the output assuming input independent normal random variables.
+        s_input = input[1].sqrt()
+        alpha = torch.clamp(-input[0] / s_input, min=-3.0, max=3.0)
+        phi = 0.5 * (1.0 + torch.erf(alpha / np.sqrt(2.0)))  # P(input<0)
+        psi = torch.exp(-0.5 * (alpha.pow(2))) / np.sqrt(2.0 * np.pi)
+        ev_gt0 = input[0] + s_input * psi / (1.0 - phi)
+        mu = (1.0 - phi) * ev_gt0
+
+        # Compute the variance of the output assuming input independent normal random variables.
+        var_gt0 = input[1] * (
+            1.0 + (alpha * psi / (1.0 - phi)) - (psi / (1.0 - phi)).pow(2)
+        )
+        var = (1 - phi) * var_gt0 + phi * (1 - phi) * ev_gt0.pow(2)
+
+        if module._module.inplace:
+            out = list(input)
+            out[0].copy_(mu)
+            out[1].copy_(var)
+            return type(input)(out)
+        else:
+            return type(input)([mu, var])
+
+
+class LeakyReLUa(MomentPropagator):
+    """Deterministic moment propagation of a normal random variable through a leaky-ReLU.
+
+    NOTE: the suffix "a" is added to prevent fastbnns.bnn.wrappers.covert_to_bnn_() from
+    automatically selecting this propagator for LeakyReLU. This is currently desirable
+    since using UnscentedTransform is faster (though less accurate) than analytical propagation.
+    """
+
+    def __init__(self):
+        """Initializer for leaky-ReLU inference module"""
+        super().__init__()
+
+    def forward(
+        self,
+        module: BayesianModule,
+        input: Iterable[torch.Tensor],
+    ) -> Iterable[torch.Tensor]:
+        """Analytical moment propagation through layer."""
+        ## Compute analytical result under mean-field approximation following
+        ## https://doi.org/10.48550/arXiv.2402.14532
+        # Compute the mean of the output assuming input independent normal random variables.
+        l = -module._module.negative_slope
+        s_input = input[1].sqrt()
+        alpha = torch.clamp(-input[0] / s_input, min=-3.0, max=3.0)
+        phi = 0.5 * (1.0 + torch.erf(alpha / np.sqrt(2.0)))  # P(input<0)
+        psi = torch.exp(-0.5 * (alpha.pow(2))) / np.sqrt(2.0 * np.pi)
+        ev_lt0 = input[0] - s_input * psi / phi
+        ev_gt0 = input[0] + s_input * psi / (1.0 - phi)
+        mu = l * phi * ev_lt0 + (1.0 - phi) * ev_gt0
+
+        # Compute the variance of the output assuming input independent normal random variables.
+        var_lt0 = input[1] * (1.0 - (alpha * psi / phi) - (psi / phi).pow(2))
+        var_gt0 = input[1] * (
+            1.0 + (alpha * psi / (1.0 - phi)) - (psi / (1.0 - phi)).pow(2)
+        )
+        var = (
+            (l**2) * phi * var_lt0
+            + (1 - phi) * var_gt0
+            + phi * (1 - phi) * (l * ev_lt0 - ev_gt0).pow(2)
+        )
+
+        if module._module.inplace:
+            out = list(input)
+            out[0].copy_(mu)
+            out[1].copy_(var)
+            return type(input)(out)
+        else:
+            return type(input)([mu, var])
 
 
 if __name__ == "__main__":
     """Example usages of inference modules."""
     import matplotlib.pyplot as plt
+    from fastbnns.bnn.wrappers import BayesianModule
 
     # Define a nonlinearity to propagate through.
-    layer = torch.nn.LeakyReLU()
+    layer = BayesianModule(module=torch.nn.LeakyReLU())
 
     # Define some propagators.
     mc = MonteCarlo(n_samples=100)
     ut = UnscentedTransform()
+    lr = LeakyReLUa()
 
     # Propagate example data through layer.
     input = (torch.tensor([1.23])[None, :], torch.tensor([3.21])[None, :])
     out_mc, samples_mc = mc(module=layer, input=input, return_samples=True)
     out_ut, samples_ut = ut(module=layer, input=input, return_samples=True)
+    out_lr = lr(module=layer, input=input)
 
     # Plot results.
     x = torch.linspace(
@@ -541,5 +716,14 @@ def __init__(self):
         .exp(),
         label="Unscented transform estimated PDF",
     )
+    ax.plot(
+        x,
+        torch.distributions.Normal(
+            loc=out_lr[0].squeeze(), scale=out_lr[1].squeeze().sqrt()
+        )
+        .log_prob(x)
+        .exp(),
+        label="Analytical PDF",
+    )
     plt.legend()
     plt.show()
diff --git a/src/fastbnns/bnn/wrappers.py b/src/fastbnns/bnn/wrappers.py
index 2b51781..9d1a676 100644
--- a/src/fastbnns/bnn/wrappers.py
+++ b/src/fastbnns/bnn/wrappers.py
@@ -91,9 +91,12 @@ def isolate_leaf_module_names(module_names: list[str]) -> list[str]:
 def convert_to_bnn_(
     model: torch.nn.Module,
     layer_wrappers: dict = {},
+    layer_wrappers_tag: dict = {},
+    layer_wrappers_type: dict = {},
     wrapper_kwargs: dict = {},
+    wrapper_kwargs_tag: dict = {},
+    wrapper_kwargs_type: dict = {},
     wrapper_kwargs_global: dict = {},
-    broadcast_module_tags: Union[list, tuple] = (),
 ) -> None:
     """Convert layers of `model` to Bayesian counterparts.
 
@@ -103,56 +106,101 @@ def convert_to_bnn_(
             module layers. The keys are names of leaf modules (e.g.,
             "module1.layer1") and the values are the corresponding wrapper
             class present in this module (e.g., "BroadcastModule").
+        layer_wrappers_tag: Extends functionality of `layer_wrappers` where
+            keys don't have to be exact layer names but instead can be tags,
+            e.g., {"encoder": "BroadcastModule"} specifies that any
+            submodule in [name for name, _ in model.named_modules()] satisfying
+            `"encoder" in name` will be wrapped with a BroadcastModule.
+        layer_wrappers_type: Extends functionality of `layer_wrappers` where
+            keys are now type names, e.g., {"BatchNorm2d": "BroadcastModule"}
+            specifies that any submodule in [m for m in model.modules()]
+            satisfying `"BatchNorm2d" == type(m).__name__"` will be wrapped
+            with a BroadcastModule.
         wrapper_kwargs: Additional keyword arguments passed to
             initialization of named Bayesian layers.  For example, if `model`
             has a module named "module1", we'll convert "module1" as
             Converter(module1, **wrapper_kwargs["module1"]) where
             Converter is a module converter.
+        wrapper_kwargs_tag: Extends functionality of `wrapper_kwargs` where
+            keys don't have to be exact layer names but instead can be tags,
+            e.g., {"encoder": encoder_kwargs} specifies that any
+            submodule in [name for name, _ in model.named_modules()] satisfying
+            `"encoder" in name` will use `encoder_kwargs` for their wrapper.
+        wrapper_kwargs_type: Extends functionality of `wrapper_kwargs` where
+            keys are now type names, e.g., {"BatchNorm2d": batchnorm_kwargs}
+            specifies that any submodule in [m for m in model.modules()]
+            satisfying `"BatchNorm2d" == type(m).__name__"` will use
+            `batchnorm_kwargs` for their wrapper.
         wrapper_kwargs_global: Keyword arguments that we'll merge
-            with values of wrapper_kwargs as, e.g.,
-            Converter(module1, **(wrapper_kwargs_global | wrapper_kwargs["module1"]))
-        broadcast_module_tags: List of strings that, if present in the class
-            name of a module, will indicate the module should be treated as a
-            broadcast module, i.e., apply forward method to mean and variance
-            directly without additional logic.
+            with values of wrapper_kwargs, wrapper_kwargs_tag, and
+            wrapper_kwargs_type as appropriate for each module.
     """
     # Search for modules of `model` to convert, removing stem modules from the
     # list (we just want the leaf modules that contain parameters).
-    module_names = [n for n, _ in model.named_modules()]
-    leaf_names = isolate_leaf_module_names(module_names)
+    modules = {k: v for k, v in model.named_modules()}
+    leaf_names = isolate_leaf_module_names(list(modules.keys()))
 
     # Replace leaf modules with Bayesian counterparts or compatible passthroughs.
     for leaf in leaf_names:
-        module = model.get_submodule(leaf)
-        module_name = module.__class__.__name__
-
-        # Prepare module arguments.
-        module_kwargs = wrapper_kwargs_global | wrapper_kwargs.pop(leaf, {})
-
-        # Search for an appropriate module converter, in the following order of
-        # priority:
-        #   (1) User-specified wrapper designated in `layer_wrappers`.
-        #   (2) Broadcast layer if tagged by broadcast_module_tags or listed
-        #       in PASSTHROUGH list.
-        #   (3) Named converters if a wrapper exists with the same name as the
-        #       module class.
-        #   (4) BayesianModule
-        if wrapper := layer_wrappers.pop(leaf, None):
+        module = modules[leaf]
+        module_name = type(module).__name__
+
+        # Prepare module arguments in the following order of priority:
+        #   (1) User-specified kwargs in `wrapper_kwargs`
+        #   (2) User-specified kwargs in `wrapper_kwargs_tag`
+        #   (3) User-specified kwargs in `wrapper_kwargs_type`
+        #   (4) Global default kwargs in `wrapper_kwargs_global`
+        if custom_kwargs := wrapper_kwargs.pop(leaf, {}):
+            pass
+        elif custom_kwargs := [v for k, v in wrapper_kwargs_tag.items() if k in leaf]:
+            # If there are multiple matches we'll just use the first one.
+            custom_kwargs = custom_kwargs[0]
+        elif custom_kwargs := [
+            v for k, v in wrapper_kwargs_type.items() if module_name == k
+        ]:
+            # If there are multiple matches we'll just use the first one.
+            custom_kwargs = custom_kwargs[0]
+        else:
+            custom_kwargs = {}
+        module_kwargs = wrapper_kwargs_global | custom_kwargs
+
+        ## Search for an appropriate module converter, in the following order of
+        ## priority:
+        #   (1) User-specified wrapper designated in `layer_wrappers`
+        #   (2) User-specified wrapper designated in `layer_wrappers_tag`
+        #   (3) User-specified wrapper designated in `layer_wrappers_type`
+        #   (4) Named converters if a wrapper exists with the same name as the
+        #       module class
+        #   (5) BroadcastModule if the module class is listed in BROADCAST
+        #   (6) BayesianModule
+
+        if custom_wrapper := layer_wrappers.pop(leaf, None):
+            pass
+        elif custom_wrapper := [v for k, v in layer_wrappers_tag.items() if k in leaf]:
+            # If there are multiple matches we'll just use the first one.
+            custom_wrapper = custom_wrapper[0]
+        elif custom_wrapper := [
+            v for k, v in layer_wrappers_type.items() if module_name == k
+        ]:
+            # If there are multiple matches we'll just use the first one.
+            custom_wrapper = custom_wrapper[0]
+        else:
+            custom_wrapper = None
+
+        if custom_wrapper is not None:
             # Wrap module in user-specified wrapper.
-            bayesian_layer = getattr(CURRENT_MODULE, wrapper)(
+            bayesian_layer = getattr(CURRENT_MODULE, custom_wrapper)(
                 module=module, **module_kwargs
             )
-        elif (module_name in BROADCAST) or any(
-            [tag in module_name for tag in broadcast_module_tags]
-        ):
-            # This module can be broadcast along (mu, var) without additional
-            # processing (e.g., a flatten layer, which only changes shapes).
-            bayesian_layer = BroadcastModule(module=module, **module_kwargs)
         elif hasattr(CURRENT_MODULE, module_name):
             # If a custom converter exists for this named layer, we'll use that by default.
             bayesian_layer = getattr(CURRENT_MODULE, module_name)(
                 module=module, **module_kwargs
             )
+        elif module_name in BROADCAST:
+            # This module can be broadcast along (mu, var) without additional
+            # processing (e.g., a flatten layer, which only changes shapes).
+            bayesian_layer = BroadcastModule(module=module, **module_kwargs)
         else:
             bayesian_layer = BayesianModule(module=module, **module_kwargs)
 
diff --git a/src/fastbnns/models/lightning_wrappers.py b/src/fastbnns/models/lightning_wrappers.py
new file mode 100644
index 0000000..30fd2bc
--- /dev/null
+++ b/src/fastbnns/models/lightning_wrappers.py
@@ -0,0 +1,66 @@
+"""PyTorch Lightning modules for BNNs."""
+
+from typing import Any, Callable
+from functools import partial
+
+import lightning as L
+import torch
+
+from fastbnns.bnn import base, losses, types
+
+
+class BNNLightning(L.LightningModule):
+    """PyTorch Lightning wrapper for BNN class."""
+
+    def __init__(
+        self,
+        bnn: base.BNN,
+        loss: losses.BNNLoss,
+        optimizer: Callable = partial(torch.optim.AdamW, lr=1.0e-3),
+    ) -> None:
+        """Initialize Lightning wrapper.
+
+        Args:
+            bnn: Bayesian neural network to wrap in Lightning.
+            loss: Loss function to call in training/validation.
+            optimizer: Partially initialized optimizer that will be given parameters
+                to optimize in self.configure_optimizers().
+        """
+        super().__init__()
+
+        self.bnn = bnn
+        self.loss = loss
+        self.optimizer_fxn = optimizer
+
+    def forward(self, *args, **kwargs) -> Any:
+        """Forward pass through BNN."""
+        return self.bnn(*args, **kwargs)
+
+    def configure_optimizers(self):
+        return self.optimizer_fxn(self.parameters())
+
+    def training_step(self, batch, batch_idx):
+        """Training step for a single batch."""
+        # Compute forward pass through model.
+        out = self.bnn(types.MuVar(batch[0]))
+
+        # Compute loss.
+        loss = self.loss(model=self.bnn, input=out[0], target=batch[1], var=out[1])
+
+        # Log results.
+        self.log("train_loss", loss, prog_bar=True, sync_dist=True)
+
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        """Validation step for a single batch."""
+        # Compute forward pass through model.
+        out = self.bnn(types.MuVar(batch[0]))
+
+        # Compute loss.
+        loss = self.loss(model=self.bnn, input=out[0], target=batch[1], var=out[1])
+
+        # Log results.
+        self.log("val_loss", loss, prog_bar=True, sync_dist=True)
+
+        return loss
diff --git a/tests/bnn/test_inference.py b/tests/bnn/test_inference.py
index 2871b70..a562351 100644
--- a/tests/bnn/test_inference.py
+++ b/tests/bnn/test_inference.py
@@ -24,7 +24,7 @@ def test_inference() -> None:
     in_features = 3
     x = MuVar(
         torch.randn((batch_size, in_features)),
-        torch.zeros((batch_size, in_features)),
+        torch.randn((batch_size, in_features)) ** 2,
     )
     propagators = [
         BasicPropagator(),
@@ -45,7 +45,7 @@ def test_inference() -> None:
     in_features = 1
     x = MuVar(
         torch.randn((batch_size, in_features)),
-        torch.zeros((batch_size, in_features)),
+        torch.randn((batch_size, in_features)) ** 2,
     )
     n_samples = 100
     propagators = [
@@ -80,7 +80,7 @@ def test_inference() -> None:
     out_features = 2
     x = MuVar(
         torch.randn((batch_size, in_features)),
-        torch.zeros((batch_size, in_features)),
+        torch.randn((batch_size, in_features)) ** 2,
     )
     module = torch.nn.Linear(in_features=in_features, out_features=out_features)
     bayes_module = BayesianModule(module, learn_var=True)
@@ -100,9 +100,10 @@ def test_inference() -> None:
             torch.randn(
                 (batch_size, in_features, *[kernel_size for _ in range(n_dim[n])])
             ),
-            torch.zeros(
+            torch.randn(
                 (batch_size, in_features, *[kernel_size for _ in range(n_dim[n])])
-            ),
+            )
+            ** 2,
         )
         module = getattr(torch.nn, f"Conv{n_dim[n]}d")(
             in_channels=in_features,
@@ -126,9 +127,10 @@ def test_inference() -> None:
             torch.randn(
                 (batch_size, in_features, *[kernel_size for _ in range(n_dim[n])])
             ),
-            torch.zeros(
+            torch.randn(
                 (batch_size, in_features, *[kernel_size for _ in range(n_dim[n])])
-            ),
+            )
+            ** 2,
         )
         module = getattr(torch.nn, f"ConvTranspose{n_dim[n]}d")(
             in_channels=in_features,