Kernel Gaussian

gaussian

CLASS	DESCRIPTION
GaussianKernel	Gaussian kernel for computing pairwise similarity matrices.

FUNCTION	DESCRIPTION
gaussian_kernel	Compute Gaussian kernel matrix from distance matrix (functional interface).

Classes

GaussianKernelResult

Bases: NamedTuple

Result container for Gaussian kernels.

ATTRIBUTE	DESCRIPTION
kernel	Computed Gaussian kernel matrix. TYPE: Tensor
bandwidth	Squared bandwidth matrix used in computation, broadcasted to match kernel shape. TYPE: Tensor

GaussianKernel(bw_method: str | torch.Tensor = 'median', learnable: bool = False, learnable_log: bool = True, normalization: KernelNormalizer | None = None, bw_kwargs: dict | None = None, dtype: torch.dtype = torch.float32)

Bases: Kernel

Gaussian kernel for computing pairwise similarity matrices.

Implements the Gaussian kernel with flexible bandwidth estimation. The kernel transforms distances into similarities using the exponential function, with the bandwidth parameter controlling the kernel width.

Supports both square (n_samples, n_samples) and non-square (n_samples, m_samples) distance matrices. Non-square matrices are useful for out-of-sample extensions like Nyström approximation.

Kernel form: \(K(d_{kl}) = \exp(-d_{kl}^2 / \sigma_{kl}^2)\), where \(d_{kl}\) is pairwise distance between samples \(k\) and \(l\), and \(\sigma_{kl}^2\) is the squared bandwidth (can be constant or adaptive per sample pair).

PARAMETER	DESCRIPTION
bw_method	Bandwidth estimation method or constant value. Available methods: "median": Median of all pairwise distances "median_1d": Per-dimension median bandwidths "quantile": Specified quantile of distances "quantile_1d": Per-dimension quantile bandwidths "quantile_search": Adaptive quantile search "knn": K-nearest neighbor based bandwidth "bgh": BGH method [1] torch.Tensor: Constant bandwidth (scalar or broadcastable tensor) TYPE: str | torch.Tensor, optional, by default "median" DEFAULT: 'median'
learnable	Whether constant bandwidth should be learnable during training. TYPE: bool, optional, by default False DEFAULT: False
learnable_log	Whether to parameterize learnable bandwidth in log space for numerical stability. TYPE: bool, optional, by default True DEFAULT: True
normalization	Kernel normalization instance. If None, no normalization is applied. TYPE: KernelNormalizer | None, optional, by default None DEFAULT: None
bw_kwargs	Parameters passed to bandwidth computation. See bw_compute in spectre.kernel.bw for complete list. TYPE: dict | None, optional, by default None DEFAULT: None
dtype	Data type for tensors. TYPE: torch.dtype, optional, by default torch.float32 DEFAULT: float32

ATTRIBUTE	DESCRIPTION
bw	Squared bandwidth matrix \(\sigma^2\) with shape matching the input distance matrix. Always stores the squared, broadcasted bandwidth for consistent computation across all bandwidth methods. Set after each forward pass. TYPE: Tensor
bw_method	The bandwidth estimation method being used. TYPE: str
bw_constant	Constant bandwidth value if bw_method="const", otherwise None. TYPE: Tensor | None
bw_learnable	Learnable bandwidth parameter if learnable=True, otherwise None. TYPE: Parameter | None

Examples:

Basic usage with median bandwidth:

>>> import torch
>>> from spectre.kernel import GaussianKernel
>>> distances = torch.rand(100, 100)
>>> kernel = GaussianKernel(bw_method="median")
>>> K = kernel(distances)
>>> K.shape
torch.Size([100, 100])

Constant bandwidth:

>>> kernel = GaussianKernel(bw_method=torch.tensor(0.5))
>>> K = kernel(distances)
>>> kernel.bw[0, 0]  # Squared bandwidth: 0.5² = 0.25
tensor(0.25)

Learnable bandwidth:

>>> kernel = GaussianKernel(
...     bw_method=torch.tensor(1.0),
...     learnable=True,
...     learnable_log=True,
... )
>>> K = kernel(distances)
>>> kernel.bw_learnable.grad  # Can optimize bandwidth via backprop

Learnable per-sample bandwidth (vector):

>>> # Learn different bandwidth for each sample (row-specific)
>>> bw_init = torch.ones(100, 1) * 0.5  # Shape (n_samples, 1)
>>> kernel = GaussianKernel(bw_method=bw_init, learnable=True)
>>> K = kernel(distances)
>>> kernel.bw_learnable.shape  # (100, 1) - one bandwidth per row

Adaptive k-NN bandwidth:

>>> kernel = GaussianKernel(bw_method="knn", bw_kwargs={"k_neighbor": 10})
>>> K = kernel(distances)

Custom quantile:

>>> kernel = GaussianKernel(bw_method="quantile", bw_kwargs={"q": 0.25})
>>> K = kernel(distances)

Notes

• Bandwidth choice affects kernel behavior:
  • Small: localized, sharp transitions
  • Large: global, smooth transitions
• Adaptive methods ("knn", "median_1d") handle varying data density better
• The bw attribute always stores σ² (squared bandwidth), broadcasted to distance matrix shape, for consistent computation across all methods
• Per-dimension methods ("median_1d", "quantile_1d") create adaptive bandwidth matrices via outer products
• Constant tensor bandwidths can be scalar or vector (broadcast to matrix shape)
• Learnable bandwidths support flexible shapes:
  • Scalar (0D or 1D with 1 element): Single bandwidth for entire kernel
  • Row vector (n, 1): Per-sample bandwidth (different for each row)
  • Column vector (1, m): Per-feature bandwidth or per-sample if the matrix represents pairwise distances (different for each column)
  • Matrix (n, m): Element-wise adaptive bandwidth (most flexible)

METHOD	DESCRIPTION
forward_step	Compute Gaussian kernel matrix.
get_kernel_params	Return current bandwidth parameter from last forward pass.
freeze_params	Freeze bandwidth to prevent recomputation on new data.
from_pdist	Create kernel with bandwidth estimated from a pairwise distance matrix.

Source code in spectre/kernel/gaussian.py

def __init__(
    self,
    bw_method: str | torch.Tensor = "median",
    learnable: bool = False,
    learnable_log: bool = True,
    normalization: KernelNormalizer | None = None,
    bw_kwargs: dict | None = None,
    dtype: torch.dtype = torch.float32,
) -> None:
    super().__init__(
        learnable=learnable,
        learnable_log=learnable_log,
        normalization=normalization,
        dtype=dtype,
    )

    if bw_kwargs is None:
        bw_kwargs = {}
    self.bw_kwargs = bw_kwargs

    if isinstance(bw_method, str):
        self.bw_method = bw_method
        self.bw_constant = None

        if bw_method in ["quantile", "quantile_1d", "quantile_search"]:
            self.q = bw_kwargs.get("q", 0.5)
            if not isinstance(self.q, (int, float)):
                raise TypeError(
                    f"Quantile `q` must be a number, got {type(self.q)}."
                )
            check_in_interval(self.q, "[0, 1]")

        if bw_method == "knn":
            self.k_neighbor = bw_kwargs.get("k_neighbor", 5)
            if not isinstance(self.k_neighbor, int):
                raise TypeError(
                    f"Number of neighbor `k_neighbor` must be an integer, got "
                    f"{type(self.k_neighbor)}."
                )
            check_in_interval(self.k_neighbor, "[1, inf)")
    else:
        # Constant bandwidth (tensor only)
        self.bw_method = "const"

        if not isinstance(bw_method, torch.Tensor):
            raise TypeError(
                f"When `bw_method` is not a string, it must be a torch.Tensor, "
                f"got {type(bw_method)}."
            )

        self.bw_constant = bw_method.to(dtype=dtype)

        if learnable:
            # Create learnable parameter with same shape as provided tensor
            if self.learnable_log:
                self.bw_learnable = torch.nn.Parameter(
                    torch.log(self.bw_constant), requires_grad=True
                )
            else:
                self.bw_learnable = torch.nn.Parameter(
                    self.bw_constant.clone(), requires_grad=True
                )

Functions

forward_step(D: torch.Tensor, weights: torch.Tensor | None = None, **kwargs) -> torch.Tensor

Compute Gaussian kernel matrix.

Delegates to gaussian_kernel() function and caches bandwidth.

PARAMETER	DESCRIPTION
D	Distance matrix, shape (n_samples, n_samples) or (n_samples, m_samples). TYPE: Tensor
weights	Per-sample importance weights for reweighted bandwidth estimation. When provided, bandwidth is estimated under the reweighted distribution, correcting for sampling bias in biased simulations. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
**kwargs	Additional keyword arguments passed to the kernel. TYPE: dict DEFAULT: {}

RETURNS	DESCRIPTION
Tensor	Gaussian kernel matrix with same shape as D.

Notes

• Delegates all computation to gaussian_kernel() function
• Sets GaussianKernel.bw to the squared, broadcasted bandwidth matrix
• Normalization (if required) is handled by parent Kernel.forward()
• When weights are provided and bw_method is a string (e.g., "median"), the bandwidth estimation accounts for importance weights to prevent inflated bandwidths from biased sampling

Source code in spectre/kernel/gaussian.py

def forward_step(
    self, D: torch.Tensor, weights: torch.Tensor | None = None, **kwargs
) -> torch.Tensor:
    """
    Compute Gaussian kernel matrix.

    Delegates to `gaussian_kernel()` function and caches bandwidth.

    Parameters
    ----------
    D : torch.Tensor
        Distance matrix, shape (n_samples, n_samples) or (n_samples, m_samples).

    weights : torch.Tensor | None, optional, by default None
        Per-sample importance weights for reweighted bandwidth estimation.
        When provided, bandwidth is estimated under the reweighted
        distribution, correcting for sampling bias in biased simulations.

    **kwargs : dict
        Additional keyword arguments passed to the kernel.

    Returns
    -------
    torch.Tensor
        Gaussian kernel matrix with same shape as D.

    Notes
    -----
    - Delegates all computation to `gaussian_kernel()` function
    - Sets `GaussianKernel.bw` to the squared, broadcasted bandwidth matrix
    - Normalization (if required) is handled by parent `Kernel.forward()`
    - When `weights` are provided and `bw_method` is a string (e.g.,
      "median"), the bandwidth estimation accounts for importance weights
      to prevent inflated bandwidths from biased sampling
    """
    if self.learnable:
        # Pass learnable parameter directly as bw_method
        bw_method_arg = self._get_param(self.bw_learnable)
    else:
        bw_method_arg = (
            self.bw_method if self.bw_method != "const" else self.bw_constant
        )

    result = gaussian_kernel(
        D,
        bw_method=bw_method_arg,
        bw_kwargs=self.bw_kwargs,
        dtype=self.dtype,
        weights=weights,
    )

    self.bw = result.bandwidth

    return result.kernel

get_kernel_params() -> dict[str, torch.Tensor]

Return current bandwidth parameter from last forward pass.

Returns the actual bandwidth matrix used in kernel computation. This is \(\sigma^2\) (squared bandwidth) broadcast to match the distance matrix shape from the most recent forward pass.

For compact scalar statistics, use get_param_summary() instead.

RETURNS	DESCRIPTION
dict[str, Tensor]	Dictionary with "bandwidth" key containing \(\sigma^2\) matrix from last forward pass. Returns empty dict if forward() has not been called yet.

Examples:

Basic usage after forward pass:

>>> kernel = GaussianKernel(bw_method=torch.tensor(1.0), learnable=True)
>>> D = torch.rand(3, 3)
>>> _ = kernel(D)
>>> params = kernel.get_kernel_params()
>>> params["bandwidth"].shape
torch.Size([3, 3])
>>> # Returns squared, broadcast bandwidth used in computation

Before forward pass:

>>> kernel = GaussianKernel(bw_method="median")
>>> params = kernel.get_kernel_params()
>>> params
{}

Logging kernel parameters during training with PyTorch Lightning:

>>> import torch
>>> from spectre.parametric import SpectralMap
>>> from spectre.kernel import GaussianKernel
>>> from spectre.utils.callbacks import MetricsCallback, kernel_params_extractor
>>> from pytorch_lightning import Trainer
>>>
>>> # Create spectral map with learnable bandwidth
>>> kernel = GaussianKernel(bw_method=torch.tensor(1.0), learnable=True)
>>> spectral_map = SpectralMap(model=my_encoder, kernel_fn=kernel, n_states=3)
>>>
>>> # Create callback to log kernel parameters
>>> callback = MetricsCallback(
...     extractors=[kernel_params_extractor],
...     log_every_n_epochs=1,
... )
>>>
>>> # Train with callback
>>> trainer = Trainer(callbacks=[callback], max_epochs=100)
>>> trainer.fit(spectral_map, train_dataloader)
>>>
>>> # Logged metrics will include:
>>> # - bandwidth_mean: mean of squared bandwidth matrix
>>> # - bandwidth_std: std of squared bandwidth matrix
>>> # - bandwidth_min: min of squared bandwidth matrix
>>> # - bandwidth_max: max of squared bandwidth matrix
>>> # - kernel_type: "GaussianKernel"
>>> # - kernel_learnable: True

Notes

The returned bandwidth is \(\sigma^2\), not \(\sigma\), and is broadcast to the shape of the distance matrix from the last forward pass. This represents the actual values used in the kernel computation: \(K = \exp(-D^2/\sigma^2)\).

Source code in spectre/kernel/gaussian.py

def get_kernel_params(self) -> dict[str, torch.Tensor]:
    """
    Return current bandwidth parameter from last forward pass.

    Returns the actual bandwidth matrix used in kernel computation. This is
    $\\sigma^2$ (squared bandwidth) broadcast to match the distance matrix
    shape from the most recent forward pass.

    For compact scalar statistics, use `get_param_summary()` instead.

    Returns
    -------
    dict[str, torch.Tensor]
        Dictionary with "bandwidth" key containing $\\sigma^2$ matrix from
        last forward pass. Returns empty dict if `forward()` has not been
        called yet.

    Examples
    --------
    Basic usage after forward pass:

    >>> kernel = GaussianKernel(bw_method=torch.tensor(1.0), learnable=True)
    >>> D = torch.rand(3, 3)
    >>> _ = kernel(D)
    >>> params = kernel.get_kernel_params()
    >>> params["bandwidth"].shape
    torch.Size([3, 3])
    >>> # Returns squared, broadcast bandwidth used in computation

    Before forward pass:

    >>> kernel = GaussianKernel(bw_method="median")
    >>> params = kernel.get_kernel_params()
    >>> params
    {}

    Logging kernel parameters during training with PyTorch Lightning:

    >>> import torch
    >>> from spectre.parametric import SpectralMap
    >>> from spectre.kernel import GaussianKernel
    >>> from spectre.utils.callbacks import MetricsCallback, kernel_params_extractor
    >>> from pytorch_lightning import Trainer
    >>>
    >>> # Create spectral map with learnable bandwidth
    >>> kernel = GaussianKernel(bw_method=torch.tensor(1.0), learnable=True)
    >>> spectral_map = SpectralMap(model=my_encoder, kernel_fn=kernel, n_states=3)
    >>>
    >>> # Create callback to log kernel parameters
    >>> callback = MetricsCallback(
    ...     extractors=[kernel_params_extractor],
    ...     log_every_n_epochs=1,
    ... )
    >>>
    >>> # Train with callback
    >>> trainer = Trainer(callbacks=[callback], max_epochs=100)
    >>> trainer.fit(spectral_map, train_dataloader)
    >>>
    >>> # Logged metrics will include:
    >>> # - bandwidth_mean: mean of squared bandwidth matrix
    >>> # - bandwidth_std: std of squared bandwidth matrix
    >>> # - bandwidth_min: min of squared bandwidth matrix
    >>> # - bandwidth_max: max of squared bandwidth matrix
    >>> # - kernel_type: "GaussianKernel"
    >>> # - kernel_learnable: True

    Notes
    -----
    The returned bandwidth is $\\sigma^2$, not $\\sigma$, and is broadcast
    to the shape of the distance matrix from the last forward pass. This
    represents the actual values used in the kernel computation: $K =
    \\exp(-D^2/\\sigma^2)$.
    """
    if hasattr(self, "bw"):
        # Always return cached value from last forward pass
        # This is the actual bandwidth matrix used: σ² broadcast to D.shape
        return {"bandwidth": self.bw}
    else:
        # Before first forward pass, no bandwidth computed yet
        return {}

freeze_params() -> GaussianKernel

Freeze bandwidth to prevent recomputation on new data.

Converts dynamic bandwidth methods (e.g., "median", "scott") to constant bandwidth using the current cached value from the last forward pass.

RETURNS	DESCRIPTION
GaussianKernel	Returns self for method chaining.

RAISES	DESCRIPTION
RuntimeError	If called before any forward pass (no bandwidth cached).

Source code in spectre/kernel/gaussian.py

def freeze_params(self) -> "GaussianKernel":
    """
    Freeze bandwidth to prevent recomputation on new data.

    Converts dynamic bandwidth methods (e.g., "median", "scott") to
    constant bandwidth using the current cached value from the last forward
    pass.

    Returns
    -------
    GaussianKernel
        Returns self for method chaining.

    Raises
    ------
    RuntimeError
        If called before any forward pass (no bandwidth cached).
    """
    if self.bw_method == "const":
        return self

    if not hasattr(self, "bw") or self.bw is None:
        raise RuntimeError(
            "Cannot freeze parameters before computing them. "
            "Call the kernel on data first to compute bandwidth."
        )

    self.bw_method = "const"
    self.bw_constant = torch.sqrt(self.bw.mean()).detach()

    return self

from_pdist(D: torch.Tensor, bw_method: str = 'median', learnable: bool = False, learnable_log: bool = True, normalization: KernelNormalizer | None = None, bw_kwargs: dict | None = None, dtype: torch.dtype = torch.float32) -> GaussianKernel classmethod

Create kernel with bandwidth estimated from a pairwise distance matrix.

Computes bandwidth from the distance distribution using the specified method and returns a GaussianKernel with constant bandwidth. Can be used to get a data-driven estimate of bandwidth as an initial value optimized during learning.

PARAMETER	DESCRIPTION
D	Pairwise distance matrix, shape (n_samples, n_samples). TYPE: Tensor
bw_method	Bandwidth estimation method to apply to D. See bw_compute for available methods. TYPE: str, optional, by default "median" DEFAULT: 'median'
learnable	Whether the estimated bandwidth should be learnable. TYPE: bool, optional, by default False DEFAULT: False
learnable_log	Whether to parameterize in log space when learnable. TYPE: bool, optional, by default True DEFAULT: True
normalization	Kernel normalization specification. TYPE: KernelNormalizer | None, optional, by default None DEFAULT: None
bw_kwargs	Additional parameters for bandwidth estimation. TYPE: dict | None, optional, by default None DEFAULT: None
dtype	Data type. TYPE: torch.dtype, optional, by default torch.float32 DEFAULT: float32

RETURNS	DESCRIPTION
GaussianKernel	Kernel with constant bandwidth derived from D.

Examples:

>>> import torch
>>> from spectre.kernel import GaussianKernel
>>> X = torch.randn(100, 10)
>>> D = torch.cdist(X, X)
>>> kernel = GaussianKernel.from_pdist(D, bw_method="median")
>>> K = kernel(D)

Source code in spectre/kernel/gaussian.py

@classmethod
def from_pdist(
    cls,
    D: torch.Tensor,
    bw_method: str = "median",
    learnable: bool = False,
    learnable_log: bool = True,
    normalization: KernelNormalizer | None = None,
    bw_kwargs: dict | None = None,
    dtype: torch.dtype = torch.float32,
) -> "GaussianKernel":
    """
    Create kernel with bandwidth estimated from a pairwise distance matrix.

    Computes bandwidth from the distance distribution using the specified
    method and returns a `GaussianKernel` with constant bandwidth. Can be
    used to get a data-driven estimate of bandwidth as an initial value
    optimized during learning.

    Parameters
    ----------
    D : torch.Tensor
        Pairwise distance matrix, shape `(n_samples, n_samples)`.

    bw_method : str, optional, by default "median"
        Bandwidth estimation method to apply to `D`. See `bw_compute`
        for available methods.

    learnable : bool, optional, by default False
        Whether the estimated bandwidth should be learnable.

    learnable_log : bool, optional, by default True
        Whether to parameterize in log space when learnable.

    normalization : KernelNormalizer | None, optional, by default None
        Kernel normalization specification.

    bw_kwargs : dict | None, optional, by default None
        Additional parameters for bandwidth estimation.

    dtype : torch.dtype, optional, by default torch.float32
        Data type.

    Returns
    -------
    GaussianKernel
        Kernel with constant bandwidth derived from `D`.

    Examples
    --------
    >>> import torch
    >>> from spectre.kernel import GaussianKernel
    >>> X = torch.randn(100, 10)
    >>> D = torch.cdist(X, X)
    >>> kernel = GaussianKernel.from_pdist(D, bw_method="median")
    >>> K = kernel(D)
    """
    check_2d(D)

    if bw_kwargs is None:
        bw_kwargs = {}

    bw = torch.sqrt(
        bw_compute(D, bw_method=bw_method, bw_kwargs=bw_kwargs).mean()
    ).to(dtype=dtype)

    return cls(
        bw_method=bw,
        learnable=learnable,
        learnable_log=learnable_log,
        normalization=normalization,
        dtype=dtype,
    )

Functions

gaussian_kernel(D: torch.Tensor, bw_method: str | torch.Tensor = 'median', bw_kwargs: dict | None = None, dtype: torch.dtype = torch.float32, weights: torch.Tensor | None = None) -> GaussianKernelResult

Compute Gaussian kernel matrix from distance matrix (functional interface).

Core Gaussian kernel computation. Returns both the kernel matrix and the bandwidth values used in computation.

This is a pure function that computes the unnormalized Gaussian kernel. For normalized kernels, use the GaussianKernel class with a normalization.

PARAMETER	DESCRIPTION
D	Distance matrix, shape (n_samples, n_samples) or (n_samples, m_samples). TYPE: Tensor
bw_method	Bandwidth estimation method or constant value. Available methods: "median": Median of all pairwise distances "median_1d": Per-dimension median bandwidths "quantile": Specified quantile of distances "quantile_1d": Per-dimension quantile bandwidths "quantile_search": Adaptive quantile search "knn": K-nearest neighbor based bandwidth "bgh": BGH method [1] torch.Tensor: Constant or learnable bandwidth (scalar or broadcastable tensor). Learnable tensors (with requires_grad=True) are detected automatically. TYPE: str | torch.Tensor, optional, by default "median" DEFAULT: 'median'
bw_kwargs	Additional parameters for bandwidth methods via dict. Supported keys: "q": float, quantile value for "quantile*" methods (0 to 1, default 0.5) "k_neighbor": int, number of neighbors for "knn" method (default 5) TYPE: dict | None, optional, by default None DEFAULT: None
dtype	Data type for tensors. TYPE: torch.dtype, optional, by default torch.float32 DEFAULT: float32
weights	Per-sample importance weights, shape (n_samples,). When provided, bandwidth is estimated under the reweighted distribution to correct for sampling bias. Only used for string-based bandwidth methods; ignored for constant/learnable tensor bandwidths. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
GaussianKernelResult	NamedTuple containing: kernel: Unnormalized Gaussian kernel matrix with same shape as D bandwidth: Squared bandwidth matrix used in computation

Examples:

Basic usage with median bandwidth:

>>> import torch
>>> from spectre.kernel import gaussian_kernel
>>> D = torch.rand(50, 50)
>>> result = gaussian_kernel(D, bw_method="median")
>>> result.kernel.shape
torch.Size([50, 50])
>>> result.bandwidth.shape
torch.Size([50, 50])

Constant bandwidth:

>>> result = gaussian_kernel(D, bw_method=torch.tensor(0.5))
>>> result.kernel[0, 0]  # Diagonal should be 1.0
tensor(1.)

Custom quantile:

>>> result = gaussian_kernel(D, bw_method="quantile", bw_kwargs={"q": 0.25})

Adaptive k-NN bandwidth:

>>> result = gaussian_kernel(D, bw_method="knn", bw_kwargs={"k_neighbor": 10})

Source code in spectre/kernel/gaussian.py

def gaussian_kernel(
    D: torch.Tensor,
    bw_method: str | torch.Tensor = "median",
    bw_kwargs: dict | None = None,
    dtype: torch.dtype = torch.float32,
    weights: torch.Tensor | None = None,
) -> GaussianKernelResult:
    """
    Compute Gaussian kernel matrix from distance matrix (functional interface).

    Core Gaussian kernel computation. Returns both the kernel matrix and the
    bandwidth values used in computation.

    This is a pure function that computes the unnormalized Gaussian kernel. For
    normalized kernels, use the `GaussianKernel` class with a normalization.

    Parameters
    ----------
    D : torch.Tensor
        Distance matrix, shape (n_samples, n_samples) or (n_samples, m_samples).

    bw_method : str | torch.Tensor, optional, by default "median"
        Bandwidth estimation method or constant value.

        Available methods:

        - "median": Median of all pairwise distances
        - "median_1d": Per-dimension median bandwidths
        - "quantile": Specified quantile of distances
        - "quantile_1d": Per-dimension quantile bandwidths
        - "quantile_search": Adaptive quantile search
        - "knn": K-nearest neighbor based bandwidth
        - "bgh": BGH method [@berry2015nonparametric]
        - torch.Tensor: Constant or learnable bandwidth (scalar or
          broadcastable tensor). Learnable tensors (with `requires_grad=True`)
          are detected automatically.

    bw_kwargs : dict | None, optional, by default None
        Additional parameters for bandwidth methods via dict.

        Supported keys:

        - "q": float, quantile value for "quantile*" methods (0 to 1, default
          0.5)
        - "k_neighbor": int, number of neighbors for "knn" method (default 5)

    dtype : torch.dtype, optional, by default torch.float32
        Data type for tensors.

    weights : torch.Tensor | None, optional, by default None
        Per-sample importance weights, shape (n_samples,). When provided,
        bandwidth is estimated under the reweighted distribution to correct for
        sampling bias. Only used for string-based bandwidth methods; ignored
        for constant/learnable tensor bandwidths.

    Returns
    -------
    GaussianKernelResult
        NamedTuple containing:

        - kernel: Unnormalized Gaussian kernel matrix with same shape as D
        - bandwidth: Squared bandwidth matrix used in computation

    Examples
    --------
    Basic usage with median bandwidth:

    >>> import torch
    >>> from spectre.kernel import gaussian_kernel
    >>> D = torch.rand(50, 50)
    >>> result = gaussian_kernel(D, bw_method="median")
    >>> result.kernel.shape
    torch.Size([50, 50])
    >>> result.bandwidth.shape
    torch.Size([50, 50])

    Constant bandwidth:

    >>> result = gaussian_kernel(D, bw_method=torch.tensor(0.5))
    >>> result.kernel[0, 0]  # Diagonal should be 1.0
    tensor(1.)

    Custom quantile:

    >>> result = gaussian_kernel(D, bw_method="quantile", bw_kwargs={"q": 0.25})

    Adaptive k-NN bandwidth:

    >>> result = gaussian_kernel(D, bw_method="knn", bw_kwargs={"k_neighbor": 10})
    """
    check_2d(D)

    if bw_kwargs is None:
        bw_kwargs = {}

    if isinstance(bw_method, str):
        bw = bw_compute(D, bw_method=bw_method, bw_kwargs=bw_kwargs, weights=weights)
    elif isinstance(bw_method, torch.Tensor):
        # Tensor bandwidth (learnable or constant)
        # For learnable tensors, preserve requires_grad by not converting dtype
        if bw_method.requires_grad:
            bw = bw_method**2
        else:
            bw = bw_method.to(dtype=dtype, device=D.device) ** 2
    else:
        raise TypeError(
            f"`bw_method` must be a string or torch.Tensor, got {type(bw_method)}."
        )

    # Broadcast bandwidth to matrix shape (n_samples, m_samples)
    bw_broadcast = safe_broadcast(bw, D, param_name="bandwidth")

    kernel = torch.exp(-torch.pow(D, 2) / bw_broadcast)

    return GaussianKernelResult(kernel=kernel, bandwidth=bw_broadcast)