polystep 0.7.0__py3-none-any.whl

polystep/layers/__init__.py

@@ -0,0 +1,16 @@
+"""Vmap-compatible PyTorch layers (attention, LSTM).
+
+- VmapSafeMultiHeadAttention: avoids SDPA (PyTorch Issue #151558)
+- VmapSafeLSTMCell / VmapSafeLSTM: explicit gate ops bypassing CuDNN
+
+Trades some throughput for guaranteed vmap compatibility.
+"""
+
+from .attention import VmapSafeMultiHeadAttention
+from .rnn import VmapSafeLSTMCell, VmapSafeLSTM
+
+__all__ = [
+    "VmapSafeMultiHeadAttention",
+    "VmapSafeLSTMCell",
+    "VmapSafeLSTM",
+]

polystep/layers/attention.py

@@ -0,0 +1,243 @@
+"""Vmap-compatible multi-head attention.
+
+A drop-in replacement for ``nn.MultiheadAttention`` that uses an
+explicit ``torch.matmul`` path instead of
+``F.scaled_dot_product_attention``. The hand-rolled path sidesteps the
+SDPA mask-validation issues reported under ``torch.vmap`` (PyTorch
+issue #151558 and related) and works the same way on every PyTorch
+build supported by this project.
+
+Example:
+    >>> import torch
+    >>> from polystep.layers import VmapSafeMultiHeadAttention
+    >>> attn = VmapSafeMultiHeadAttention(embed_dim=64, num_heads=4)
+    >>> x = torch.randn(2, 10, 64)  # (batch, seq, embed)
+    >>> out = attn(x, x, x)  # self-attention
+    >>> out.shape
+    torch.Size([2, 10, 64])
+"""
+
+import math
+import warnings
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class VmapSafeMultiHeadAttention(nn.Module):
+    """Multi-head attention using explicit matmul operations for vmap compatibility.
+
+    This implementation avoids F.scaled_dot_product_attention which has known
+    vmap mask validation bugs (PyTorch Issue #151558). Instead, it uses explicit
+    matrix multiplications that work correctly under torch.vmap.
+
+    Limitations vs nn.MultiheadAttention:
+        - No built-in causal masking (pass attn_mask manually)
+        - Assumes batch-first layout: (batch, seq, embed_dim)
+        - No add_bias_kv or add_zero_attn support
+        - kdim/vdim must equal embed_dim (no cross-attention with different dims)
+
+    Args:
+        embed_dim: Total dimension of the model (must be divisible by num_heads).
+        num_heads: Number of parallel attention heads.
+        dropout: Dropout probability on attention weights. Default: 0.0.
+        bias: Whether to add bias to projection layers. Default: True.
+
+    Input shapes:
+        - query: ``(batch, seq_q, embed_dim)``
+        - key: ``(batch, seq_k, embed_dim)``
+        - value: ``(batch, seq_k, embed_dim)``
+        - attn_mask: ``(seq_q, seq_k)``, ``(batch, seq_q, seq_k)``, or
+          ``(batch, num_heads, seq_q, seq_k)`` (float additive or bool)
+        - key_padding_mask: ``(batch, seq_k)`` bool, ``True`` = padding
+
+    Output shape:
+        - (batch, seq_q, embed_dim)
+
+    Example:
+        >>> attn = VmapSafeMultiHeadAttention(embed_dim=64, num_heads=4)
+        >>> x = torch.randn(2, 10, 64)
+        >>> out = attn(x, x, x)  # self-attention
+        >>> out.shape
+        torch.Size([2, 10, 64])
+    """
+
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        dropout: float = 0.0,
+        bias: bool = True,
+        # Mirror the upstream nn.MultiheadAttention signature so that
+        # callers passing unsupported features get a clear
+        # NotImplementedError instead of a generic Python
+        # "got an unexpected keyword argument".
+        add_bias_kv: bool = False,
+        add_zero_attn: bool = False,
+        kdim: Optional[int] = None,
+        vdim: Optional[int] = None,
+        batch_first: bool = True,
+    ):
+        super().__init__()
+
+        if embed_dim % num_heads != 0:
+            raise ValueError(f"embed_dim ({embed_dim}) must be divisible by num_heads ({num_heads})")
+
+        if kdim is not None and kdim != embed_dim:
+            raise NotImplementedError(
+                f"VmapSafeMultiHeadAttention only supports kdim == embed_dim, "
+                f"got kdim={kdim}, embed_dim={embed_dim}. See LIMITATIONS.md."
+            )
+        if vdim is not None and vdim != embed_dim:
+            raise NotImplementedError(
+                f"VmapSafeMultiHeadAttention only supports vdim == embed_dim, "
+                f"got vdim={vdim}, embed_dim={embed_dim}. See LIMITATIONS.md."
+            )
+        if add_bias_kv:
+            raise NotImplementedError(
+                "VmapSafeMultiHeadAttention does not support add_bias_kv=True. See LIMITATIONS.md."
+            )
+        if add_zero_attn:
+            raise NotImplementedError(
+                "VmapSafeMultiHeadAttention does not support add_zero_attn=True. See LIMITATIONS.md."
+            )
+        if not batch_first:
+            raise NotImplementedError(
+                "VmapSafeMultiHeadAttention only supports batch_first=True "
+                "(input layout (batch, seq, embed_dim)). See LIMITATIONS.md."
+            )
+
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        self.scale = 1.0 / math.sqrt(self.head_dim)
+        self.dropout = dropout
+
+        if dropout > 0:
+            warnings.warn(
+                "VmapSafeMultiHeadAttention with dropout > 0 requires eval mode under vmap. "
+                "Call model.eval() before vmap evaluation.",
+                stacklevel=2,
+            )
+
+        # Separate projections for Q, K, V (not combined like GPT-2)
+        # This makes vmap over parameters cleaner
+        self.W_q = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.W_k = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.W_v = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+        # Output projection
+        self.W_o = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+        # Dropout for attention weights
+        self.attn_dropout = nn.Dropout(dropout)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+        key_padding_mask: Optional[torch.Tensor] = None,
+        need_weights: bool = False,
+        is_causal: bool = False,
+    ) -> torch.Tensor:
+        """Compute multi-head attention.
+
+        Args:
+            query: Query tensor of shape ``(batch, seq_q, embed_dim)``.
+            key: Key tensor of shape ``(batch, seq_k, embed_dim)``.
+            value: Value tensor of shape ``(batch, seq_k, embed_dim)``.
+            attn_mask: Attention mask. Accepted shapes (matching
+                ``nn.MultiheadAttention``):
+                ``(seq_q, seq_k)``, ``(batch, seq_q, seq_k)``, or
+                ``(batch, num_heads, seq_q, seq_k)``. Float masks are added
+                to attention scores before softmax (use ``-inf`` for hard
+                masking). Bool masks are interpreted upstream-style: ``True``
+                positions are masked out.
+            key_padding_mask: Boolean mask of shape ``(batch, seq_k)`` where
+                ``True`` marks padding positions.
+            need_weights: Not supported. Returning attention weights from
+                inside ``vmap`` requires extra reshapes that defeat the
+                kernel fusion this layer is here to enable.
+            is_causal: Not supported. Pass an explicit triangular
+                ``attn_mask`` instead.
+
+        Returns:
+            Output tensor of shape (batch, seq_q, embed_dim).
+        """
+        # Reject unsupported forward kwargs up front so the failure
+        # mode is loud, not "wrong but plausible-looking output".
+        if need_weights:
+            raise NotImplementedError(
+                "VmapSafeMultiHeadAttention does not support need_weights=True. See LIMITATIONS.md."
+            )
+        if is_causal:
+            raise NotImplementedError(
+                "VmapSafeMultiHeadAttention does not support is_causal=True. "
+                "Pass an explicit triangular attn_mask instead. See LIMITATIONS.md."
+            )
+
+        batch_size, seq_q, _ = query.shape
+        seq_k = key.shape[1]
+
+        # Project Q, K, V: (batch, seq, embed_dim) -> (batch, seq, embed_dim)
+        Q = self.W_q(query)
+        K = self.W_k(key)
+        V = self.W_v(value)
+
+        # Reshape to (batch, num_heads, seq, head_dim)
+        Q = Q.view(batch_size, seq_q, self.num_heads, self.head_dim).transpose(1, 2)
+        K = K.view(batch_size, seq_k, self.num_heads, self.head_dim).transpose(1, 2)
+        V = V.view(batch_size, seq_k, self.num_heads, self.head_dim).transpose(1, 2)
+
+        # Compute attention scores: (batch, num_heads, seq_q, seq_k)
+        # Q @ K.T = (batch, num_heads, seq_q, head_dim) @ (batch, num_heads, head_dim, seq_k)
+        scores = torch.matmul(Q, K.transpose(-2, -1)) * self.scale
+
+        # Bool masks are mask-filled with -inf to match
+        # ``nn.MultiheadAttention`` semantics. Float masks are additive
+        # (use -inf for hard masking, finite values for soft biases like
+        # ALiBi or relative-position embeddings).
+        if attn_mask is not None:
+            if attn_mask.dim() == 2:
+                # (seq_q, seq_k) -> (1, 1, seq_q, seq_k)
+                attn_mask = attn_mask.unsqueeze(0).unsqueeze(0)
+            elif attn_mask.dim() == 3:
+                # (batch, seq_q, seq_k) -> (batch, 1, seq_q, seq_k)
+                attn_mask = attn_mask.unsqueeze(1)
+            elif attn_mask.dim() != 4:
+                raise ValueError(
+                    "attn_mask must have 2, 3, or 4 dimensions, got "
+                    f"{attn_mask.dim()}D (shape {tuple(attn_mask.shape)})"
+                )
+            if attn_mask.dtype == torch.bool:
+                scores = scores.masked_fill(attn_mask, float("-inf"))
+            else:
+                scores = scores + attn_mask
+
+        # Apply key padding mask
+        if key_padding_mask is not None:
+            # (batch, seq_k) -> (batch, 1, 1, seq_k)
+            padding_mask = key_padding_mask.unsqueeze(1).unsqueeze(2)
+            # True means padding, so we mask with -inf
+            scores = scores.masked_fill(padding_mask, float("-inf"))
+
+        # Softmax and dropout
+        attn_weights = F.softmax(scores, dim=-1)
+        attn_weights = self.attn_dropout(attn_weights)
+
+        # Apply attention to values: (batch, num_heads, seq_q, head_dim)
+        context = torch.matmul(attn_weights, V)
+
+        # Reshape back: (batch, num_heads, seq_q, head_dim) -> (batch, seq_q, embed_dim)
+        # Use reshape instead of .contiguous().view() to avoid a full tensor copy
+        # per candidate under vmap (N candidates x this copy = ~1.3 GB overhead).
+        context = context.transpose(1, 2).reshape(batch_size, seq_q, self.embed_dim)
+
+        # Output projection
+        output = self.W_o(context)
+
+        return output

polystep/layers/rnn.py

@@ -0,0 +1,276 @@
+"""Vmap-compatible LSTM bypassing CuDNN.
+
+Standard ``nn.LSTM`` (via ``aten::lstm.input``) has no batching rule under
+``torch.vmap`` and fails on both CPU and CUDA -- still the case on PyTorch
+2.12 (verified). This implementation uses explicit per-step gate
+computations, which are vmap-safe but ~2-5x slower than CuDNN.
+
+Example:
+    >>> import torch
+    >>> from polystep.layers import VmapSafeLSTM
+    >>> lstm = VmapSafeLSTM(input_size=32, hidden_size=64, num_layers=2)
+    >>> x = torch.randn(4, 10, 32)  # (batch, seq, input)
+    >>> out, (h_n, c_n) = lstm(x)
+    >>> out.shape
+    torch.Size([4, 10, 64])
+"""
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+
+class VmapSafeLSTMCell(nn.Module):
+    """LSTM cell using explicit gate computations for vmap compatibility.
+
+    This implementation avoids CuDNN which has known issues under torch.vmap
+    with ".data access" errors. Instead, explicit gate computations are used
+    that work correctly under vectorization.
+
+    Args:
+        input_size: The number of expected features in the input.
+        hidden_size: The number of features in the hidden state.
+        bias: Whether to add bias to the linear transformations. Default: True.
+
+    Input shapes:
+        - x: (batch, input_size)
+        - state: Tuple of (h, c) where both are (batch, hidden_size)
+
+    Output shapes:
+        - h_new: (batch, hidden_size) - new hidden state
+        - (h_new, c_new): Tuple of new states for compatibility with nn.LSTMCell
+
+    Example:
+        >>> cell = VmapSafeLSTMCell(input_size=32, hidden_size=64)
+        >>> x = torch.randn(4, 32)
+        >>> h = torch.zeros(4, 64)
+        >>> c = torch.zeros(4, 64)
+        >>> h_new, (h_out, c_out) = cell(x, (h, c))
+    """
+
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        bias: bool = True,
+    ):
+        super().__init__()
+
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+
+        # Combined input and hidden projections for all 4 gates
+        # Gates: input (i), forget (f), cell input (g), output (o)
+        self.W_i = nn.Linear(input_size, 4 * hidden_size, bias=bias)
+        self.W_h = nn.Linear(hidden_size, 4 * hidden_size, bias=bias)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        state: Tuple[torch.Tensor, torch.Tensor],
+    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        """Compute one LSTM step.
+
+        Args:
+            x: Input tensor of shape (batch, input_size).
+            state: Tuple of (h, c) hidden and cell states,
+                each of shape (batch, hidden_size).
+
+        Returns:
+            Tuple of:
+                - h_new: New hidden state of shape (batch, hidden_size)
+                - (h_new, c_new): Tuple of new states
+        """
+        h, c = state
+
+        # Compute all gates at once: (batch, 4 * hidden_size)
+        gates = self.W_i(x) + self.W_h(h)
+
+        # Split into 4 gates: each (batch, hidden_size)
+        i, f, g, o = gates.chunk(4, dim=-1)
+
+        # Apply activations
+        i = torch.sigmoid(i)  # Input gate
+        f = torch.sigmoid(f)  # Forget gate
+        g = torch.tanh(g)  # Cell input (candidate)
+        o = torch.sigmoid(o)  # Output gate
+
+        # Update cell state
+        c_new = f * c + i * g
+
+        # Compute new hidden state
+        h_new = o * torch.tanh(c_new)
+
+        return h_new, (h_new, c_new)
+
+
+class VmapSafeLSTM(nn.Module):
+    """Multi-layer LSTM using explicit gate computations for vmap compatibility.
+
+    This implementation wraps multiple VmapSafeLSTMCell layers to provide
+    a drop-in replacement for nn.LSTM that works under torch.vmap.
+
+    Limitations vs nn.LSTM:
+        - No bidirectional support
+        - No proj_size support
+        - Assumes batch-first layout: (batch, seq_len, input_size)
+        - 2-5x slower than CuDNN (explicit gate computation)
+
+    Note: Expected 2-5x slower than CuDNN but guaranteed vmap-safe.
+
+    Args:
+        input_size: The number of expected features in the input.
+        hidden_size: The number of features in the hidden state.
+        num_layers: Number of recurrent layers. Default: 1.
+        bias: Whether to add bias to the linear transformations. Default: True.
+        dropout: Dropout probability between layers (not after last layer).
+            Default: 0.0.
+
+    Input shapes:
+        - x: (batch, seq_len, input_size)
+        - state: Optional tuple of (h, c) where each is (num_layers, batch, hidden_size).
+            If None, zero-initialized states are used.
+
+    Output shapes:
+        - output: (batch, seq_len, hidden_size) - output features from last layer
+        - (h_n, c_n): Final states, each (num_layers, batch, hidden_size)
+
+    Example:
+        >>> lstm = VmapSafeLSTM(input_size=32, hidden_size=64, num_layers=2)
+        >>> x = torch.randn(4, 10, 32)  # (batch, seq, input)
+        >>> out, (h_n, c_n) = lstm(x)
+        >>> out.shape
+        torch.Size([4, 10, 64])
+    """
+
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        num_layers: int = 1,
+        bias: bool = True,
+        dropout: float = 0.0,
+        # Mirror the upstream nn.LSTM signature so that callers passing
+        # unsupported features get a clear NotImplementedError instead
+        # of a generic Python "got an unexpected keyword argument".
+        bidirectional: bool = False,
+        proj_size: int = 0,
+        batch_first: bool = True,
+    ):
+        super().__init__()
+
+        if num_layers < 1:
+            raise ValueError(f"num_layers must be >= 1, got {num_layers}")
+        if bidirectional:
+            raise NotImplementedError("VmapSafeLSTM does not support bidirectional=True. See LIMITATIONS.md.")
+        if proj_size != 0:
+            raise NotImplementedError("VmapSafeLSTM does not support proj_size != 0. See LIMITATIONS.md.")
+        if not batch_first:
+            raise NotImplementedError(
+                "VmapSafeLSTM only supports batch_first=True (input layout "
+                "(batch, seq_len, input_size)). See LIMITATIONS.md."
+            )
+
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.dropout = dropout
+
+        # Build cells for each layer
+        self.cells = nn.ModuleList()
+        for layer in range(num_layers):
+            layer_input_size = input_size if layer == 0 else hidden_size
+            self.cells.append(VmapSafeLSTMCell(layer_input_size, hidden_size, bias=bias))
+
+        # Dropout between layers (not applied after last layer)
+        self.dropout_layer = nn.Dropout(dropout) if dropout > 0 else None
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        state: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        """Process sequence through multi-layer LSTM.
+
+        Args:
+            x: Input tensor of shape (batch, seq_len, input_size).
+            state: Optional tuple of (h, c) initial states, each of shape
+                (num_layers, batch, hidden_size). If None, zero states are used.
+
+        Returns:
+            Tuple of:
+                - output: Tensor of shape (batch, seq_len, hidden_size)
+                    containing output features from the last layer
+                - (h_n, c_n): Tuple of final states, each of shape
+                    (num_layers, batch, hidden_size)
+        """
+        # PackedSequence support is not implemented. nn.LSTM accepts
+        # one, but our explicit-loop forward would silently treat it as
+        # a tensor and crash later with an opaque AttributeError. Catch
+        # it up front with a clear message.
+        if isinstance(x, nn.utils.rnn.PackedSequence):
+            raise NotImplementedError(
+                "VmapSafeLSTM does not support PackedSequence input. Pad to a dense tensor first. See LIMITATIONS.md."
+            )
+
+        batch_size, seq_len, _ = x.shape
+        device = x.device
+        dtype = x.dtype
+
+        # Initialize states if not provided
+        # Use list-based state tracking to avoid in-place updates under vmap
+        if state is None:
+            h_list = [
+                torch.zeros(batch_size, self.hidden_size, device=device, dtype=dtype) for _ in range(self.num_layers)
+            ]
+            c_list = [
+                torch.zeros(batch_size, self.hidden_size, device=device, dtype=dtype) for _ in range(self.num_layers)
+            ]
+        else:
+            h, c = state
+            # Split stacked tensor into list of per-layer states
+            h_list = [h[i] for i in range(self.num_layers)]
+            c_list = [c[i] for i in range(self.num_layers)]
+
+        # Process sequence timestep by timestep
+        outputs = []
+        for t in range(seq_len):
+            # Get input for this timestep: (batch, input_size)
+            x_t = x[:, t, :]
+
+            # Process through layers, collecting new states
+            new_h_list = []
+            new_c_list = []
+            for layer_idx, cell in enumerate(self.cells):
+                h_layer = h_list[layer_idx]
+                c_layer = c_list[layer_idx]
+
+                # Cell forward pass
+                h_new, (h_out, c_out) = cell(x_t, (h_layer, c_layer))
+
+                # Collect new states (no in-place update)
+                new_h_list.append(h_out)
+                new_c_list.append(c_out)
+
+                # Apply dropout between layers (not after last)
+                if self.dropout_layer is not None and layer_idx < self.num_layers - 1:
+                    x_t = self.dropout_layer(h_new)
+                else:
+                    x_t = h_new
+
+            # Update state lists for next timestep
+            h_list = new_h_list
+            c_list = new_c_list
+
+            # Collect output from last layer
+            outputs.append(x_t)
+
+        # Stack outputs: (batch, seq_len, hidden_size)
+        output = torch.stack(outputs, dim=1)
+
+        # Stack final states: (num_layers, batch, hidden_size)
+        h_n = torch.stack(h_list, dim=0)
+        c_n = torch.stack(c_list, dim=0)
+
+        return output, (h_n, c_n)

polystep/objectives/__init__.py

@@ -0,0 +1,25 @@
+from .base import ObjectiveFn
+from .synthetic import (
+    Ackley,
+    Rosenbrock,
+    Rastrigin,
+    StyblinskiTang,
+    Levy,
+    Griewank,
+    Beale,
+    Branin,
+    Sphere,
+)
+
+__all__ = [
+    "ObjectiveFn",
+    "Ackley",
+    "Rosenbrock",
+    "Rastrigin",
+    "StyblinskiTang",
+    "Levy",
+    "Griewank",
+    "Beale",
+    "Branin",
+    "Sphere",
+]

polystep/objectives/base.py

@@ -0,0 +1,77 @@
+"""Abstract base class for optimization objective functions."""
+
+import abc
+from typing import Optional
+
+import torch
+
+
+class ObjectiveFn(abc.ABC):
+    """Base class for optimization objectives.
+
+    Attributes:
+        dim: Dimensionality of the problem.
+        bounds: Search space bounds of shape ``(dim, 2)`` as
+            ``[(min, max), ...]``.
+        optimizers: Known global optimizer locations of shape
+            ``(num_opts, dim)``.
+        optimal_value: Known global optimum value.
+        noise_std: Standard deviation of additive Gaussian observation noise.
+            ``None`` or ``0`` disables noise.
+        negate: If ``True``, negate the output (turn a maximization problem
+            into a minimization problem).
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        bounds: Optional[torch.Tensor] = None,
+        optimizers: Optional[torch.Tensor] = None,
+        optimal_value: Optional[float] = None,
+        noise_std: Optional[float] = None,
+        negate: bool = False,
+    ):
+        self.dim = dim
+        self.bounds = bounds
+        self.optimizers = optimizers
+        self.optimal_value = optimal_value
+        self.noise_std = noise_std
+        self.negate = negate
+
+    @abc.abstractmethod
+    def evaluate(self, X: torch.Tensor) -> torch.Tensor:
+        """Compute the raw objective value.
+
+        Args:
+            X: Input array of shape ``(..., dim)``.
+
+        Returns:
+            Cost array of shape ``(...)``.
+        """
+        pass
+
+    def __call__(
+        self,
+        X: torch.Tensor,
+        generator: Optional[torch.Generator] = None,
+    ) -> torch.Tensor:
+        """Compute the final cost, applying noise and negation if configured.
+
+        Args:
+            X: Input points of shape ``(..., dim)``.
+            generator: Optional ``torch.Generator`` for reproducible noise.
+
+        Returns:
+            Cost values of shape ``(...)``.
+        """
+        cost = self.evaluate(X)
+        if self.noise_std is not None and self.noise_std > 0.0:
+            noise = torch.empty_like(cost).normal_(
+                mean=0.0,
+                std=float(self.noise_std),
+                generator=generator,
+            )
+            cost = cost + noise
+        if self.negate:
+            return -cost
+        return cost