noxton 0.1.0__tar.gz

noxton-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Artur A. Galstyan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

noxton-0.1.0/PKG-INFO

@@ -0,0 +1,11 @@
+Metadata-Version: 2.4
+Name: noxton
+Version: 0.1.0
+Summary: Add your description here
+Requires-Python: >=3.13
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: equinox>=0.13.4
+Requires-Dist: jax>=0.9.0.1
+Requires-Dist: statedict2pytree>=2.0.1
+Dynamic: license-file

noxton-0.1.0/noxton/functions/__init__.py

@@ -0,0 +1,32 @@
+from .activation import swiglu
+from .attention import (
+    create_attn_mask,
+    multi_head_attention_forward,
+    shifted_window_attention,
+)
+from .embedding import sinusoidal_embedding
+from .masking import (
+    build_attention_mask,
+    canonical_attn_mask,
+    canonical_key_padding_mask,
+    canonical_mask,
+    make_causal_mask,
+)
+from .normalization import normalize
+from .regularization import dropout, stochastic_depth
+
+__all__ = [
+    "swiglu",
+    "multi_head_attention_forward",
+    "shifted_window_attention",
+    "create_attn_mask",
+    "sinusoidal_embedding",
+    "dropout",
+    "stochastic_depth",
+    "normalize",
+    "build_attention_mask",
+    "canonical_attn_mask",
+    "canonical_key_padding_mask",
+    "canonical_mask",
+    "make_causal_mask",
+]

noxton-0.1.0/noxton/functions/activation.py

@@ -0,0 +1,34 @@
+import jax
+import jax.numpy as jnp
+from jaxtyping import Array
+
+
+def swiglu(x: Array, axis: int = -1) -> Array:
+    """Apply the SwiGLU activation function.
+
+    Splits the input array into two halves along the specified axis, applies
+    the Swish activation to the first half, and multiplies it element-wise
+    with the second half. This gated activation is commonly used in
+    transformer feed-forward blocks.
+
+    Args:
+        x: Input array. Its size along ``axis`` must be even.
+        axis: Axis along which to split the input into two halves.
+            Defaults to ``-1`` (last axis).
+
+    Returns:
+        Array of the same dtype as ``x`` with size halved along ``axis``.
+
+    Example:
+        >>> import jax.numpy as jnp
+        >>> x = jnp.array([1.0, 2.0, -1.0, 0.5])
+        >>> swiglu(x)  # splits into [1., 2.] and [-1., 0.5]
+        Array([-0.26894143,  0.9526741 ], dtype=float32)
+
+        >>> # 2-D input, split along last axis (default)
+        >>> x2d = jnp.ones((3, 4))
+        >>> swiglu(x2d).shape
+        (3, 2)
+    """
+    a, b = jnp.split(x, 2, axis=axis)
+    return jax.nn.swish(a) * b

noxton-0.1.0/noxton/functions/attention.py

@@ -0,0 +1,594 @@
+import functools
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+from beartype.typing import Any
+from jaxtyping import Array, Bool, Float, PRNGKeyArray
+
+from noxton.functions.normalization import normalize
+from noxton.functions.regularization import dropout as dropout_fn
+from noxton.utils.utils import default_floating_dtype
+
+
+def multi_head_attention_forward(
+    query: Float[Array, "tgt_len d_model"],
+    key: Float[Array, "src_len d_model"],
+    value: Float[Array, "src_len d_model"],
+    embed_dim_to_check: int,
+    num_heads: int,
+    in_proj_weight: Float[Array, "3*d_model d_model"] | None = None,
+    in_proj_bias: Float[Array, "3*d_model"] | None = None,
+    bias_k: Float[Array, "1 d_model"] | None = None,
+    bias_v: Float[Array, "1 d_model"] | None = None,
+    add_zero_attn: bool = False,
+    dropout_p: float = 0.0,
+    out_proj_weight: Float[Array, "d_model d_model"] | None = None,
+    out_proj_bias: Float[Array, "d_model"] | None = None,
+    inference: bool = False,
+    key_padding_mask: Float[Array, "src_len"] | Bool[Array, "src_len"] | None = None,
+    attn_mask: Float[Array, "tgt_len src_len"] | None = None,
+    need_weights: bool = True,
+    use_separate_proj_weight: bool = False,
+    q_proj_weight: Float[Array, "d_model d_model"] | None = None,
+    k_proj_weight: Float[Array, "d_model d_model"] | None = None,
+    v_proj_weight: Float[Array, "d_model d_model"] | None = None,
+    static_k: Float[Array, "src_len d_model"] | None = None,
+    static_v: Float[Array, "src_len d_model"] | None = None,
+    average_attn_weights: bool = True,
+    is_causal: bool = False,
+    dropout_key: PRNGKeyArray | None = None,
+) -> tuple[
+    Float[Array, "tgt_len d_model"],
+    Float[Array, "num_heads tgt_len src_len"]
+    | Float[Array, "tgt_len src_len"]
+    | Float[Array, "tgt_len src_len+1"]
+    | None,
+]:
+    """Compute scaled dot-product multi-head attention.
+
+    This is the functional core of multi-head attention as described in
+    *Attention Is All You Need* (Vaswani et al., 2017). It accepts
+    pre-allocated weight matrices and returns the attended output along with
+    optional attention weights.
+
+    This function is a 1:1 mapping from the PyTorch implementation.
+
+    The function supports several advanced options:
+
+    - **Separate projection weights** (``use_separate_proj_weight=True``):
+      pass individual ``q_proj_weight``, ``k_proj_weight``, and
+      ``v_proj_weight`` instead of a single fused ``in_proj_weight``.
+    - **Static key/value**: bypass the key/value projections by providing
+      pre-computed ``static_k`` / ``static_v`` tensors.
+    - **Extra key/value bias tokens** (``bias_k``, ``bias_v``): append
+      learnable bias tokens to the key and value sequences.
+    - **Zero-attention slot** (``add_zero_attn``): append a zero vector to
+      keys and values, giving the model an option to attend to nothing.
+    - **Causal masking** (``is_causal``): apply an upper-triangular ``-inf``
+      mask so each query can only attend to earlier or equal positions.
+    - **Key-padding mask** / **additive attention mask**: arbitrary masking
+      via ``key_padding_mask`` and ``attn_mask``.
+
+    Args:
+        query: Query tensor of shape ``(tgt_len, d_model)``.
+        key: Key tensor of shape ``(src_len, d_model)``.
+        value: Value tensor of shape ``(src_len, d_model)``.
+        embed_dim_to_check: Expected embedding dimension; asserted to equal
+            ``d_model``.
+        num_heads: Number of attention heads. Must divide ``d_model`` evenly.
+        in_proj_weight: Fused input projection weight of shape
+            ``(3*d_model, d_model)``. Required when
+            ``use_separate_proj_weight=False``.
+        in_proj_bias: Optional bias for the fused projection of shape
+            ``(3*d_model,)``.
+        bias_k: Optional learnable key bias of shape ``(1, d_model)`` appended
+            to the key sequence.
+        bias_v: Optional learnable value bias of shape ``(1, d_model)``
+            appended to the value sequence. Must be provided together with
+            ``bias_k``.
+        add_zero_attn: If ``True``, append a zero token to keys and values.
+            Defaults to ``False``.
+        dropout_p: Dropout probability applied to attention weights during
+            training. Defaults to ``0.0``.
+        out_proj_weight: Output projection weight of shape
+            ``(d_model, d_model)``. Required.
+        out_proj_bias: Optional output projection bias of shape ``(d_model,)``.
+        inference: If ``True``, disable dropout (eval mode).
+            Defaults to ``False``.
+        key_padding_mask: Optional mask of shape ``(src_len,)`` marking padded
+            key positions. ``True`` / non-zero values are masked out (set to
+            ``-inf`` in logits).
+        attn_mask: Optional additive attention mask of shape
+            ``(tgt_len, src_len)``. Values are added to attention logits
+            before softmax.
+        need_weights: If ``True``, also return the attention weight matrix.
+            Defaults to ``True``.
+        use_separate_proj_weight: If ``True``, use ``q_proj_weight``,
+            ``k_proj_weight``, ``v_proj_weight`` for projections instead of
+            ``in_proj_weight``. Defaults to ``False``.
+        q_proj_weight: Query projection weight ``(d_model, d_model)``.
+            Required when ``use_separate_proj_weight=True``.
+        k_proj_weight: Key projection weight ``(d_model, d_model)``.
+            Required when ``use_separate_proj_weight=True``.
+        v_proj_weight: Value projection weight ``(d_model, d_model)``.
+            Required when ``use_separate_proj_weight=True``.
+        static_k: Pre-computed key of shape ``(src_len, d_model)``. When
+            provided, bypasses the key projection.
+        static_v: Pre-computed value of shape ``(src_len, d_model)``. When
+            provided, bypasses the value projection.
+        average_attn_weights: If ``True``, average attention weights over
+            heads before returning. Defaults to ``True``.
+        is_causal: If ``True``, apply a causal mask to prevent attending to
+            future positions. Defaults to ``False``.
+        dropout_key: JAX PRNG key for attention dropout. Required when
+            ``dropout_p > 0.0`` and ``inference=False``.
+
+    Returns:
+        A tuple ``(attn_output, attn_weights)`` where:
+
+        - ``attn_output``: Attended output of shape ``(tgt_len, d_model)``.
+        - ``attn_weights``: Attention weights or ``None`` when
+          ``need_weights=False``.  Shape depends on ``average_attn_weights``:
+          ``(tgt_len, src_len)`` when averaged, or
+          ``(num_heads, tgt_len, src_len)`` otherwise.
+
+    Example:
+        >>> import jax
+        >>> import jax.numpy as jnp
+        >>> key = jax.random.PRNGKey(0)
+        >>> d_model, num_heads, tgt_len, src_len = 8, 2, 4, 6
+        >>> q = jax.random.normal(key, (tgt_len, d_model))
+        >>> k = jax.random.normal(key, (src_len, d_model))
+        >>> v = jax.random.normal(key, (src_len, d_model))
+        >>> W = jax.random.normal(key, (3 * d_model, d_model))
+        >>> W_out = jax.random.normal(key, (d_model, d_model))
+        >>> out, weights = multi_head_attention_forward(
+        ...     q, k, v,
+        ...     embed_dim_to_check=d_model,
+        ...     num_heads=num_heads,
+        ...     in_proj_weight=W,
+        ...     out_proj_weight=W_out,
+        ...     inference=True,
+        ... )
+        >>> out.shape
+        (4, 8)
+        >>> weights.shape  # averaged over heads by default
+        (4, 6)
+    """
+    tgt_len, d_model = query.shape
+    src_len, k_dim = key.shape
+    value_len, v_dim = value.shape
+
+    assert d_model == k_dim == v_dim == embed_dim_to_check, (
+        "Embedding dimensions must match"
+    )
+
+    assert src_len == value_len, "Key and value must have the same sequence length"
+
+    head_dim = d_model // num_heads
+    assert head_dim * num_heads == d_model, "embed_dim must be divisible by num_heads"
+
+    if dropout_p > 0.0:
+        assert dropout_key is not None, (
+            "dropout_key must be provided if dropout_p > 0.0"
+        )
+
+    if use_separate_proj_weight:
+        # When using separate projection weights for q, k, v
+        assert q_proj_weight is not None, (
+            "q_proj_weight should not be None when use_separate_proj_weight=True"
+        )
+        assert k_proj_weight is not None, (
+            "k_proj_weight should not be None when use_separate_proj_weight=True"
+        )
+        assert v_proj_weight is not None, (
+            "v_proj_weight should not be None when use_separate_proj_weight=True"
+        )
+
+        q = query @ q_proj_weight.T
+
+        if static_k is None:
+            k = key @ k_proj_weight.T
+        else:
+            k = static_k
+            src_len, _ = k.shape
+
+        if static_v is None:
+            v = value @ v_proj_weight.T
+        else:
+            v = static_v
+            value_len, _ = v.shape
+
+        if in_proj_bias is not None:
+            q_bias, k_bias, v_bias = jnp.split(in_proj_bias, 3)
+            q = q + q_bias
+            k = k + k_bias
+            v = v + v_bias
+
+    else:
+        assert in_proj_weight is not None, (
+            "in_proj_weight should not be None when use_separate_proj_weight=False"
+        )
+
+        q_proj_weight_part, k_proj_weight_part, v_proj_weight_part = jnp.split(
+            in_proj_weight, 3
+        )
+
+        q = query @ q_proj_weight_part.T
+
+        if static_k is None:
+            k = key @ k_proj_weight_part.T
+        else:
+            k = static_k
+            src_len, _ = static_k.shape
+
+        if static_v is None:
+            v = value @ v_proj_weight_part.T
+        else:
+            v = static_v
+            value_len, _ = static_v.shape
+
+        if in_proj_bias is not None:
+            q_bias, k_bias, v_bias = jnp.split(in_proj_bias, 3)
+            q = q + q_bias
+            k = k + k_bias
+            v = v + v_bias
+
+    assert src_len == value_len
+
+    q = q.reshape(tgt_len, num_heads, head_dim)
+    k = k.reshape(src_len, num_heads, head_dim)
+    v = v.reshape(src_len, num_heads, head_dim)
+
+    if add_zero_attn:
+        zero_attn_shape = (1, num_heads, head_dim)
+        k_zeros = jnp.zeros(zero_attn_shape)
+        v_zeros = jnp.zeros(zero_attn_shape)
+
+        k = jnp.concatenate([k, k_zeros], axis=0)
+        v = jnp.concatenate([v, v_zeros], axis=0)
+
+        src_len += 1
+        value_len += 1
+
+    if bias_k is not None and bias_v is not None:
+        bias_k = bias_k.reshape(1, num_heads, head_dim)
+        bias_v = bias_v.reshape(1, num_heads, head_dim)
+
+        k = jnp.concatenate([k, bias_k], axis=0)
+        v = jnp.concatenate([v, bias_v], axis=0)
+
+        src_len += 1
+        value_len += 1
+
+    assert src_len == value_len
+
+    # [tgt_len, num_heads, head_dim] → [num_heads, tgt_len, head_dim]
+    q = jnp.transpose(q, (1, 0, 2))
+
+    # [src_len, num_heads, head_dim] → [num_heads, src_len, head_dim]
+    k = jnp.transpose(k, (1, 0, 2))
+    v = jnp.transpose(v, (1, 0, 2))
+
+    scale = jnp.sqrt(head_dim)
+    attn_output_weights = jnp.matmul(q, jnp.transpose(k, (0, 2, 1))) / scale
+
+    if key_padding_mask is not None:
+        padding_mask = key_padding_mask.reshape(1, 1, src_len)
+        padding_mask = jnp.repeat(padding_mask, num_heads, axis=0)
+        padding_mask = jnp.repeat(padding_mask, tgt_len, axis=1)
+        attn_output_weights = jnp.where(
+            padding_mask, float("-inf"), attn_output_weights
+        )
+
+    if attn_mask is not None:
+        # [tgt_len, src_len] -> [num_heads, tgt_len, src_len]
+        mask = attn_mask.reshape(1, tgt_len, src_len)
+        mask = jnp.repeat(mask, num_heads, axis=0)
+        attn_output_weights = attn_output_weights + mask
+
+    if is_causal:
+        causal_mask = jnp.triu(jnp.ones((tgt_len, src_len)), k=1)
+        causal_mask = (causal_mask == 1).reshape(1, tgt_len, src_len)
+        causal_mask = jnp.repeat(causal_mask, num_heads, axis=0)
+        attn_output_weights = jnp.where(causal_mask, float("-inf"), attn_output_weights)
+
+    # [num_heads, tgt_len, src_len]
+    attn_output_weights = jax.nn.softmax(attn_output_weights, axis=-1)
+
+    if dropout_p > 0.0 and not inference:
+        assert dropout_key is not None, (
+            "dropout_key required because dropout_p > 0.0 and training"
+        )
+        dropout_mask = jax.random.bernoulli(
+            dropout_key, 1 - dropout_p, attn_output_weights.shape
+        )
+        scale = 1.0 / (1.0 - dropout_p)
+        attn_output_weights = attn_output_weights * dropout_mask * scale
+
+    attn_output = jnp.matmul(attn_output_weights, v)
+    attn_output = jnp.transpose(attn_output, (1, 0, 2))
+    attn_output = attn_output.reshape(tgt_len, d_model)
+
+    assert out_proj_weight is not None, "out_proj_weight must be provided"
+    attn_output = attn_output @ out_proj_weight.T
+
+    if out_proj_bias is not None:
+        attn_output = attn_output + out_proj_bias
+
+    if need_weights:
+        if average_attn_weights:
+            attn_output_weights = attn_output_weights.mean(axis=0)
+        return attn_output, attn_output_weights
+    else:
+        return attn_output, None
+
+
+def create_attn_mask(
+    pad_H: int,
+    pad_W: int,
+    window_size: list[int],
+    shift_size: list[int],
+    dtype: Any | None = None,
+) -> Array:
+    """Build the region-index mask used by shifted-window attention.
+
+    Assigns each spatial position in a ``(pad_H, pad_W)`` feature map an
+    integer *region index* that encodes which window it belongs to after a
+    cyclic shift.  Positions in the same window share the same region index;
+    positions in different windows have different indices.  The mask is used
+    downstream to zero out cross-window attention scores.
+
+    The region indices are computed by partitioning the height axis at
+    ``pad_H - window_size[0]`` and ``pad_H - shift_size[0]``, and the width
+    axis at ``pad_W - window_size[1]`` and ``pad_W - shift_size[1]``, then
+    assigning a unique integer to each ``(row_region, col_region)`` cell.
+
+    Args:
+        pad_H: Padded feature map height (must be a multiple of
+            ``window_size[0]``).
+        pad_W: Padded feature map width (must be a multiple of
+            ``window_size[1]``).
+        window_size: Local attention window size as ``[window_H, window_W]``.
+        shift_size: Cyclic shift amounts as ``[shift_H, shift_W]``.
+        dtype: Output array dtype. Defaults to the project's default floating
+            dtype when ``None``.
+
+    Returns:
+        Integer-region-index map of shape ``(pad_H, pad_W)`` cast to
+        ``dtype``.
+
+    Example:
+        >>> mask = create_attn_mask(8, 8, window_size=[4, 4], shift_size=[2, 2])
+        >>> mask.shape
+        (8, 8)
+    """
+    if dtype is None:
+        dtype = default_floating_dtype()
+    assert dtype is not None
+    h_boundaries = jnp.array([pad_H - window_size[0], pad_H - shift_size[0]])
+    w_boundaries = jnp.array([pad_W - window_size[1], pad_W - shift_size[1]])
+
+    h_boundaries = jnp.sort(h_boundaries)
+    w_boundaries = jnp.sort(w_boundaries)
+
+    ii, jj = jnp.indices((pad_H, pad_W))  # ii for rows, jj for columns
+
+    row_region_idx = jnp.searchsorted(h_boundaries, ii, side="right")
+    col_region_idx = jnp.searchsorted(w_boundaries, jj, side="right")
+
+    num_col_regions = len(w_boundaries) + 1
+    attn_mask = row_region_idx * num_col_regions + col_region_idx
+
+    return attn_mask.astype(dtype)
+
+
+def shifted_window_attention(
+    x: Float[Array, "H W C"],
+    qkv_weight: Float[Array, "in_dim out_dim"],
+    proj_weight: Float[Array, "out_dim out_dim"],
+    relative_position_bias: Array,
+    window_size: list[int],
+    num_heads: int,
+    shift_size: list[int],
+    attention_dropout: float = 0.0,
+    dropout: float = 0.0,
+    qkv_bias: Array | None = None,
+    proj_bias: Array | None = None,
+    logit_scale: Array | None = None,
+    inference: bool = False,
+    key: PRNGKeyArray | None = None,
+) -> Float[Array, "H W C"]:
+    """Apply Shifted-Window Multi-Head Self-Attention (Swin Attention).
+
+    Implements the window-based self-attention mechanism from the Swin
+    Transformer (Liu et al., 2021).  The input feature map is partitioned
+    into non-overlapping local windows; attention is computed independently
+    within each window.  A cyclic spatial shift (``shift_size``) is applied
+    before partitioning to create cross-window connections, and a region-index
+    mask is used to prevent attending across window boundaries introduced by
+    the padding/shift.
+
+    Supports two attention score variants:
+
+    - **Scaled dot-product** (``logit_scale=None``): standard
+      ``q @ k.T / sqrt(head_dim)``.
+    - **Cosine attention** (when ``logit_scale`` is provided): L2-normalised
+      ``q`` and ``k`` are used, and scores are multiplied by a learnable
+      (but bounded) temperature ``exp(min(logit_scale, log(100)))``.
+
+    Relative position biases are added to the attention logits before softmax.
+
+    Args:
+        x: Input feature map of shape ``(H, W, C)``.
+        qkv_weight: Fused QKV projection weight of shape
+            ``(in_dim, out_dim)`` where ``out_dim = 3 * C``.
+        proj_weight: Output projection weight of shape ``(C, C)``.
+        relative_position_bias: Relative position bias tensor added to
+            attention logits; shape must broadcast with
+            ``(num_windows, num_heads, window_H*window_W, window_H*window_W)``.
+        window_size: Local window size as ``[window_H, window_W]``.
+        num_heads: Number of attention heads. Must divide ``C`` evenly.
+        shift_size: Cyclic shift amounts as ``[shift_H, shift_W]``.
+            Use ``[0, 0]`` to disable shifting (regular window attention).
+        attention_dropout: Dropout probability applied to attention weights.
+            Defaults to ``0.0``.
+        dropout: Dropout probability applied to the output projection.
+            Defaults to ``0.0``.
+        qkv_bias: Optional bias for the QKV projection of shape
+            ``(3 * C,)``. When ``logit_scale`` is also provided, the key
+            bias component is zeroed out (as in Swin V2).
+        proj_bias: Optional bias for the output projection of shape ``(C,)``.
+        logit_scale: Learnable log-scale scalar for cosine attention. When
+            ``None``, standard scaled dot-product attention is used.
+        inference: If ``True``, disable dropout. Defaults to ``False``.
+        key: JAX PRNG key required when ``inference=False``.
+
+    Returns:
+        Output feature map of shape ``(H, W, C)``.
+
+    Raises:
+        ValueError: If ``inference=False`` and ``key`` is ``None``.
+
+    Example:
+        >>> import jax
+        >>> import jax.numpy as jnp
+        >>> H, W, C, num_heads = 8, 8, 16, 2
+        >>> window_size, shift_size = [4, 4], [2, 2]
+        >>> key = jax.random.PRNGKey(0)
+        >>> x = jax.random.normal(key, (H, W, C))
+        >>> qkv_w = jax.random.normal(key, (C, 3 * C))
+        >>> proj_w = jax.random.normal(key, (C, C))
+        >>> win_tokens = window_size[0] * window_size[1]
+        >>> num_wins = (H // window_size[0]) * (W // window_size[1])
+        >>> rpb = jnp.zeros((num_wins, num_heads, win_tokens, win_tokens))
+        >>> out = shifted_window_attention(
+        ...     x, qkv_w, proj_w, rpb,
+        ...     window_size=window_size, num_heads=num_heads,
+        ...     shift_size=shift_size, inference=True,
+        ... )
+        >>> out.shape
+        (8, 8, 16)
+    """
+    if not inference and key is None:
+        raise ValueError("Need key when in training mode")
+    H, W, C = x.shape
+    to_pad_W = (window_size[1] - W % window_size[1]) % window_size[1]
+    to_pad_H = (window_size[0] - H % window_size[0]) % window_size[0]
+    x = jnp.pad(x, ((0, to_pad_H), (0, to_pad_W), (0, 0)))
+    pad_H, pad_W, _ = x.shape
+
+    shift_size = shift_size.copy()
+    if window_size[0] >= pad_H:
+        shift_size[0] = 0
+    if window_size[1] >= pad_W:
+        shift_size[1] = 0
+
+    # cyclic shift
+    if sum(shift_size) > 0:
+        x = jnp.roll(x, shift=(-shift_size[0], -shift_size[1]), axis=(0, 1))
+
+    # partition windows
+    num_windows = (pad_H // window_size[0]) * (pad_W // window_size[1])
+    x = jnp.reshape(
+        x,
+        (
+            pad_H // window_size[0],
+            window_size[0],
+            pad_W // window_size[1],
+            window_size[1],
+            C,
+        ),
+    )
+    x = jnp.transpose(x, (0, 2, 1, 3, 4)).reshape(
+        num_windows, window_size[0] * window_size[1], C
+    )
+
+    # multi-head attention
+    if logit_scale is not None and qkv_bias is not None:
+        length = qkv_bias.size // 3
+        qkv_bias = qkv_bias.at[length : 2 * length].set(0.0)
+
+    def linear(x: Array, weight: Array, bias: Array | None):
+        output = x @ jnp.transpose(weight)  # (in,) @ (in, out) -> (out,)
+        if bias is not None:
+            output = output + bias
+        return output
+
+    linear_pt = functools.partial(linear, weight=qkv_weight, bias=qkv_bias)
+
+    qkv = eqx.filter_vmap(eqx.filter_vmap(linear_pt))(x)
+    win_size, patches, _ = qkv.shape
+    qkv = jnp.transpose(
+        qkv.reshape(win_size, patches, 3, num_heads, C // num_heads), (2, 0, 3, 1, 4)
+    )
+    q, k, v = qkv[0], qkv[1], qkv[2]
+    if logit_scale is not None:
+        # cosine attention
+        attn = normalize(q, axis=-1) @ jnp.transpose(
+            normalize(k, axis=-1), (0, 1, 3, 2)
+        )
+        # Clamp the logit scale exponent for stability
+        logit_scale = jnp.exp(jnp.minimum(logit_scale, jnp.log(jnp.array(100.0))))
+        attn = attn * logit_scale
+    else:
+        q = q * (C // num_heads) ** -0.5
+        # attn = q @ (jnp.transpose(normalize(k, axis=-1), (0, 1, 3, 2))) # Incorrect
+        attn = q @ jnp.transpose(k, (0, 1, 3, 2))  # Corrected: q @ k.T
+
+    # add relative position bias
+    attn = attn + relative_position_bias
+
+    if sum(shift_size) > 0:
+        attn_mask = create_attn_mask(pad_H, pad_W, window_size, shift_size)
+        attn_mask = attn_mask.reshape(
+            pad_H // window_size[0],
+            window_size[0],
+            pad_W // window_size[1],
+            window_size[1],
+        )
+        attn_mask = jnp.transpose(attn_mask, (0, 2, 1, 3)).reshape(
+            num_windows, window_size[0] * window_size[1]
+        )
+        attn_mask = jnp.expand_dims(attn_mask, axis=1) - jnp.expand_dims(
+            attn_mask, axis=2
+        )
+        attn_mask = jnp.where(attn_mask == 0, 0.0, -100.0)
+
+        attn = attn + attn_mask[:, None, :, :]
+
+    attn = jax.nn.softmax(attn, axis=-1)
+    if not inference:
+        assert key is not None, "key must be given if not inference"
+        key, subkey = jax.random.split(key)
+        attn = dropout_fn(attn, p=attention_dropout, inference=inference, key=subkey)
+
+    x = jnp.transpose(attn @ v, (0, 2, 1, 3)).reshape(
+        num_windows, window_size[0] * window_size[1], C
+    )
+    linear_pt_proj = functools.partial(linear, weight=proj_weight, bias=proj_bias)
+
+    x = eqx.filter_vmap(eqx.filter_vmap(linear_pt_proj))(x)
+    if not inference:
+        assert key is not None, "key must be given if not inference"
+        key, subkey = jax.random.split(key)
+        x = dropout_fn(x, p=dropout, inference=inference, key=subkey)
+
+    # reverse windows
+    x = x.reshape(
+        pad_H // window_size[0],
+        pad_W // window_size[1],
+        window_size[0],
+        window_size[1],
+        C,
+    )
+    x = jnp.transpose(x, (0, 2, 1, 3, 4)).reshape(pad_H, pad_W, C)
+
+    # reverse cyclic shift
+    if sum(shift_size) > 0:
+        x = jnp.roll(x, shift=(shift_size[0], shift_size[1]), axis=(0, 1))
+
+    # unpad features
+    x = x[:H, :W, :]
+    return x

noxton-0.1.0/noxton/functions/embedding.py

@@ -0,0 +1,64 @@
+import jax.numpy as jnp
+from beartype.typing import Any
+from jaxtyping import Array, Float, Int
+
+from noxton.utils import default_floating_dtype
+
+
+def sinusoidal_embedding(
+    t: Int[Array, ""], embedding_size: int, dtype: Any | None = None
+) -> Float[Array, " embedding_size"]:
+    """Compute a sinusoidal positional embedding for a scalar timestep.
+
+    Encodes a single integer timestep ``t`` into a fixed-size vector using
+    alternating sine and cosine functions at geometrically spaced frequencies,
+    following the scheme from *Attention Is All You Need* (Vaswani et al., 2017)
+    and commonly used in diffusion model timestep conditioning.
+
+    The embedding frequencies are::
+
+        freq_i = exp(-log(10000) * i / (embedding_size / 2))   for i in [0, half_dim)
+
+    The output is the concatenation of ``sin(t * freq)`` and ``cos(t * freq)``.
+
+    Args:
+        t: Scalar integer timestep (0-d array).
+        embedding_size: Length of the output embedding vector. Must be even.
+        dtype: Floating-point dtype for the output. Defaults to the project's
+            default floating dtype when ``None``.
+
+    Returns:
+        1-D array of shape ``(embedding_size,)`` containing the sinusoidal
+        embedding for timestep ``t``.
+
+    Raises:
+        ValueError: If ``embedding_size`` is odd.
+
+    Example:
+        >>> import jax.numpy as jnp
+        >>> t = jnp.array(100)
+        >>> emb = sinusoidal_embedding(t, embedding_size=16)
+        >>> emb.shape
+        (16,)
+
+        >>> # Embeddings for different timesteps are distinct
+        >>> emb0 = sinusoidal_embedding(jnp.array(0), embedding_size=8)
+        >>> emb1 = sinusoidal_embedding(jnp.array(1), embedding_size=8)
+        >>> bool(jnp.any(emb0 != emb1))
+        True
+    """
+    if dtype is None:
+        dtype = default_floating_dtype()
+    assert dtype is not None
+    if embedding_size % 2 != 0:
+        raise ValueError(f"Embedding size must be even, but got {embedding_size}")
+
+    half_dim = embedding_size // 2
+    embedding_freqs = jnp.exp(
+        -jnp.log(10000) * jnp.arange(start=0, stop=half_dim, dtype=dtype) / half_dim
+    )
+
+    time_args = t * embedding_freqs
+    embedding = jnp.concatenate([jnp.sin(time_args), jnp.cos(time_args)], axis=-1)
+
+    return embedding

noxton-0.1.0/noxton/functions/masking.py

@@ -0,0 +1,198 @@
+import jax.numpy as jnp
+from jaxtyping import Array, Bool
+
+
+def canonical_mask(
+    mask,
+    mask_name,
+    other_name="",
+    other_type=None,
+    target_type=jnp.float32,
+    other_mask=None,
+    check_other=True,
+):
+    """Convert an arbitrary mask tensor into a canonical additive float mask.
+
+    Accepts boolean, integer, or floating-point masks and normalises them to
+    a floating-point additive mask (``0.0`` for positions to attend to,
+    ``-inf`` for positions to ignore) that can be directly added to attention
+    logits.
+
+    - **Boolean masks**: ``True`` → ``-inf``, ``False`` → ``0.0``
+    - **Integer / float masks**: cast to ``target_type`` without modification.
+    - ``None`` input returns ``None`` (no mask).
+
+    Args:
+        mask: The mask to canonicalise. May be a boolean, integer, or
+            floating-point JAX array, or ``None``.
+        mask_name: Human-readable name of this mask used in error messages.
+        other_name: Human-readable name of the secondary mask (used in error
+            messages only). Defaults to ``""``.
+        other_type: Expected dtype of ``other_mask`` (currently unused in the
+            implementation but kept for API compatibility). Defaults to
+            ``None``.
+        target_type: Target floating-point dtype for the output. Defaults to
+            ``jnp.float32``.
+        other_mask: A secondary mask for cross-validation purposes (currently
+            unused). Defaults to ``None``.
+        check_other: Whether to perform cross-mask validation (currently
+            unused). Defaults to ``True``.
+
+    Returns:
+        A floating-point array of dtype ``target_type`` with the same shape
+        as ``mask``, or ``None`` if ``mask`` is ``None``.
+
+    Raises:
+        TypeError: If ``mask`` has an unsupported dtype (not bool, int, or
+            float).
+
+    Example:
+        >>> import jax.numpy as jnp
+        >>> bool_mask = jnp.array([True, False, True])
+        >>> canonical_mask(bool_mask, "attn_mask")
+        Array([-inf,   0., -inf], dtype=float32)
+
+        >>> float_mask = jnp.array([0.0, -1e9, 0.0])
+        >>> canonical_mask(float_mask, "attn_mask")
+        Array([ 0.e+00, -1.e+09,  0.e+00], dtype=float32)
+    """
+    if mask is None:
+        return None
+    if mask.dtype == bool:
+        additive_mask = jnp.where(mask, -jnp.inf, 0.0).astype(target_type)
+        return additive_mask
+    elif jnp.issubdtype(mask.dtype, jnp.integer) or jnp.issubdtype(
+        mask.dtype, jnp.floating
+    ):
+        return mask.astype(target_type)
+    else:
+        raise TypeError(
+            f"{mask_name} must be bool, int, or float tensor, but got {mask.dtype}"
+        )
+
+
+def canonical_key_padding_mask(
+    key_padding_mask, attn_mask=None, query_dtype=jnp.float32
+):
+    """Convert a key-padding mask to a canonical additive float mask.
+
+    A convenience wrapper around :func:`canonical_mask` that applies the
+    correct argument names for key-padding masks.  The resulting mask can be
+    directly added to attention logits: padded positions get ``-inf`` (boolean
+    ``True`` input) while real positions get ``0.0``.
+
+    Args:
+        key_padding_mask: Mask of shape ``(src_len,)`` indicating which key
+            positions are padding.  May be boolean (``True`` = padded),
+            integer, or float, or ``None``.
+        attn_mask: The attention mask that will be combined with this mask
+            (used for error reporting only). Defaults to ``None``.
+        query_dtype: Target floating-point dtype for the output. Defaults to
+            ``jnp.float32``.
+
+    Returns:
+        A floating-point array of dtype ``query_dtype`` or ``None``.
+
+    Example:
+        >>> import jax.numpy as jnp
+        >>> kpm = jnp.array([False, False, True])  # last token is padding
+        >>> canonical_key_padding_mask(kpm)
+        Array([  0.,   0., -inf], dtype=float32)
+    """
+    return canonical_mask(
+        mask=key_padding_mask,
+        mask_name="key_padding_mask",
+        other_name="attn_mask",
+        other_mask=attn_mask,
+        target_type=query_dtype,
+    )
+
+
+def canonical_attn_mask(attn_mask, query_dtype=jnp.float32):
+    """Convert an attention mask to a canonical additive float mask.
+
+    A convenience wrapper around :func:`canonical_mask` that applies the
+    correct argument names for attention masks.  Boolean masks are converted
+    so that ``True`` positions are masked out (``-inf``) and ``False``
+    positions are kept (``0.0``).  Numeric masks are cast to ``query_dtype``
+    without modification, allowing pre-built additive masks (e.g. causal
+    masks with ``0`` and ``-inf``) to be passed through unchanged.
+
+    Args:
+        attn_mask: Attention mask of shape ``(tgt_len, src_len)``.  May be
+            boolean, integer, or floating-point, or ``None``.
+        query_dtype: Target floating-point dtype for the output. Defaults to
+            ``jnp.float32``.
+
+    Returns:
+        A floating-point array of dtype ``query_dtype`` or ``None``.
+
+    Example:
+        >>> import jax.numpy as jnp
+        >>> mask = jnp.array([[False, True], [False, False]])
+        >>> canonical_attn_mask(mask)
+        Array([[  0., -inf],
+               [  0.,   0.]], dtype=float32)
+    """
+    return canonical_mask(
+        mask=attn_mask,
+        mask_name="attn_mask",
+        other_type=None,
+        other_name="",
+        target_type=query_dtype,
+        check_other=False,
+    )
+
+
+def make_causal_mask(seq_len: int) -> Bool[Array, "seq_len seq_len"]:
+    """Create a boolean lower-triangular causal mask.
+
+    Position ``[i, j]`` is ``True`` when position ``j`` is allowed to attend
+    to position ``i`` (i.e. ``j <= i``), and ``False`` otherwise.  Use this
+    mask to prevent tokens from attending to future positions.
+
+    Args:
+        seq_len: Sequence length; the output is a square matrix of shape
+            ``(seq_len, seq_len)``.
+
+    Returns:
+        Boolean array of shape ``(seq_len, seq_len)`` where the lower
+        triangle (including the diagonal) is ``True`` and the upper triangle
+        is ``False``.
+
+    Example:
+        >>> make_causal_mask(3)
+        Array([[ True, False, False],
+               [ True,  True, False],
+               [ True,  True,  True]], dtype=bool)
+    """
+    return jnp.tril(jnp.ones((seq_len, seq_len), dtype=jnp.bool_))
+
+
+def build_attention_mask(context_length: int) -> Array:
+    """Create an additive causal attention mask with ``0`` and ``-inf``.
+
+    Produces a float32 lower-triangular matrix where attended positions
+    carry ``0.0`` (no change to logits) and future positions carry ``-inf``
+    (zeroed out after softmax).  The mask can be added directly to attention
+    logit matrices.
+
+    Args:
+        context_length: Sequence length; the output is a square matrix of
+            shape ``(context_length, context_length)``.
+
+    Returns:
+        Float32 array of shape ``(context_length, context_length)`` with
+        ``0.0`` on and below the diagonal and ``-inf`` above the diagonal.
+
+    Example:
+        >>> build_attention_mask(3)
+        Array([[  0., -inf, -inf],
+               [  0.,   0., -inf],
+               [  0.,   0.,   0.]], dtype=float32)
+    """
+    mask = jnp.tril(jnp.zeros((context_length, context_length)))
+    upper = jnp.triu(jnp.full((context_length, context_length), float("-inf")), k=1)
+
+    mask = mask + upper
+    return mask

noxton-0.1.0/noxton/functions/normalization.py

@@ -0,0 +1,39 @@
+import jax.numpy as jnp
+from jaxtyping import Array
+
+
+def normalize(x: Array, p: int = 2, axis: int = 1, eps: float = 1e-12) -> Array:
+    """Normalize an array along an axis using an Lp norm.
+
+    Computes ``x / max(||x||_p, eps)`` along ``axis``, where the denominator is
+    clamped to at least ``eps`` to prevent division by zero.
+
+    Args:
+        x: Input array to normalize.
+        p: Order of the norm. ``p=2`` gives the Euclidean (L2) norm,
+            ``p=1`` gives the Manhattan (L1) norm, etc. Defaults to ``2``.
+        axis: Axis along which to compute the norm and normalize.
+            Defaults to ``1``.
+        eps: Small constant added to the denominator for numerical stability.
+            Defaults to ``1e-12``.
+
+    Returns:
+        Array of the same shape and dtype as ``x`` with unit Lp norm along
+        ``axis`` (or norm equal to ``eps`` when the input norm is smaller).
+
+    Example:
+        >>> import jax.numpy as jnp
+        >>> x = jnp.array([[3.0, 4.0], [0.0, 0.0]])
+        >>> normalize(x)          # L2 norm along axis=1
+        Array([[0.6, 0.8],
+               [0. , 0. ]], dtype=float32)
+
+        >>> normalize(x, p=1)     # L1 norm along axis=1
+        Array([[0.42857143, 0.5714286 ],
+               [0.        , 0.        ]], dtype=float32)
+    """
+    norm = jnp.linalg.norm(x, ord=p, axis=axis, keepdims=True)
+    norm = jnp.maximum(norm, eps)
+    output = x / norm
+
+    return output

noxton-0.1.0/noxton/functions/regularization.py

@@ -0,0 +1,125 @@
+import jax
+import jax.numpy as jnp
+from jaxtyping import Array, PRNGKeyArray
+
+
+def stochastic_depth(
+    input: Array,
+    p: float,
+    mode: str,
+    inference: bool,
+    key: PRNGKeyArray,
+) -> Array:
+    """Apply Stochastic Depth regularization (DropPath).
+
+    During training, randomly drops entire samples (``mode="batch"``) or
+    individual rows (``mode="row"``). At inference time or when ``p=0``, the
+    input is returned unchanged. Surviving elements are scaled by
+    ``1 / (1 - p)`` to preserve the expected value.
+
+    Args:
+        input: Input array to apply stochastic depth to.
+        p: Drop probability in ``[0, 1]``. The probability that a sample/row
+            is zeroed out. ``p=0`` disables the operation.
+        mode: Dropping granularity. One of:
+            - ``"batch"``: a single binary mask is broadcast over the whole
+              batch (all samples share the same fate).
+            - ``"row"``: each sample in the batch gets its own independent
+              mask.
+        inference: If ``True``, skip stochastic depth and return ``input``
+            unchanged (equivalent to test/eval mode).
+        key: JAX PRNG key used to sample the Bernoulli noise.
+
+    Returns:
+        Array of the same shape and dtype as ``input``.
+
+    Raises:
+        ValueError: If ``p`` is outside ``[0, 1]``.
+        ValueError: If ``mode`` is not ``"batch"`` or ``"row"``.
+
+    Example:
+        >>> import jax
+        >>> import jax.numpy as jnp
+        >>> key = jax.random.PRNGKey(0)
+        >>> x = jnp.ones((4, 8))
+        >>> # row mode: each of the 4 samples may be dropped independently
+        >>> out = stochastic_depth(x, p=0.5, mode="row", inference=False, key=key)
+        >>> out.shape
+        (4, 8)
+
+        >>> # inference mode always returns the input unchanged
+        >>> stochastic_depth(x, p=0.9, mode="batch", inference=True, key=key) is x
+        True
+    """
+    if p < 0.0 or p > 1.0:
+        raise ValueError(f"drop probability has to be between 0 and 1, but got {p}")
+    if mode not in ["batch", "row"]:
+        raise ValueError(f"mode has to be either 'batch' or 'row', but got {mode}")
+    if inference or p == 0.0:
+        return input
+    survival_rate = 1.0 - p
+    if mode == "row":
+        size = [input.shape[0]] + [1] * (input.ndim - 1)
+    else:
+        size = [1] * input.ndim
+    noise = jax.random.bernoulli(key, p=survival_rate, shape=size).astype(input.dtype)
+    if survival_rate > 0.0:
+        noise = noise / survival_rate
+    return input * noise
+
+
+def dropout(
+    x: Array,
+    p: float,
+    inference: bool,
+    key: PRNGKeyArray | None = None,
+) -> Array:
+    """Apply dropout regularization to an array.
+
+    During training, each element is independently zeroed with probability
+    ``p``. Surviving elements are scaled by ``1 / (1 - p)`` so the expected
+    sum is preserved (inverted dropout). At inference time or when ``p=0``
+    the input is returned unchanged.
+
+    Args:
+        x: Input array.
+        p: Probability of an element being zeroed. Must be in ``[0, 1)``.
+            When ``p=0`` the function is a no-op regardless of ``inference``.
+        inference: If ``True``, return ``x`` unchanged (eval/test mode).
+        key: JAX PRNG key. Required when ``inference=False`` and ``p > 0``.
+            Defaults to ``None``.
+
+    Returns:
+        Array of the same shape and dtype as ``x``.
+
+    Raises:
+        RuntimeError: If ``inference=False``, ``p > 0``, and ``key`` is
+            ``None``.
+
+    Example:
+        >>> import jax
+        >>> import jax.numpy as jnp
+        >>> key = jax.random.PRNGKey(42)
+        >>> x = jnp.ones((3, 4))
+        >>> out = dropout(x, p=0.5, inference=False, key=key)
+        >>> out.shape
+        (3, 4)
+
+        >>> # inference mode: output equals input
+        >>> import jax.numpy as jnp
+        >>> x = jnp.array([1.0, 2.0, 3.0])
+        >>> dropout(x, p=0.5, inference=True)
+        Array([1., 2., 3.], dtype=float32)
+    """
+    if isinstance(p, (int, float)) and p == 0:
+        inference = True
+    if inference:
+        return x
+    elif key is None:
+        raise RuntimeError(
+            "Dropout requires a key when running in non-deterministic mode."
+        )
+    else:
+        q = 1 - jax.lax.stop_gradient(p)
+        mask = jax.random.bernoulli(key, q, x.shape)
+        return jnp.where(mask, x / q, 0)

noxton-0.1.0/noxton/nn/abstract.py

@@ -0,0 +1,16 @@
+import abc
+
+import equinox as eqx
+from jaxtyping import Array
+
+
+class AbstractNorm(eqx.Module):
+    @abc.abstractmethod
+    def __call__(self, x: Array, *_, **__) -> Array: ...
+
+
+class AbstractNormStateful(eqx.nn.StatefulLayer):
+    @abc.abstractmethod
+    def __call__(
+        self, x: Array, state: eqx.nn.State, *_, **__
+    ) -> tuple[Array, eqx.nn.State]: ...

noxton-0.1.0/noxton/utils/__init__.py

@@ -0,0 +1,5 @@
+from .utils import default_floating_dtype
+
+__all__ = [
+    "default_floating_dtype",
+]

noxton-0.1.0/noxton/utils/utils.py

@@ -0,0 +1,59 @@
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+from beartype.typing import Any
+from jaxtyping import PyTree
+
+
+def default_floating_dtype() -> Any:
+    if jax.config.read("jax_enable_x64"):
+        return jnp.float64
+    else:
+        return jnp.float32
+
+
+def summarize_model(model: PyTree) -> str:
+    params, _ = eqx.partition(model, eqx.is_array)
+
+    param_counts = {}
+    total_params = 0
+
+    def count_params(pytree, name=""):
+        nonlocal total_params
+        count = 0
+        if isinstance(pytree, jnp.ndarray):
+            count = pytree.size
+            total_params += count
+            if name:
+                param_counts[name] = count
+        elif hasattr(pytree, "__dict__"):
+            for key, value in pytree.__dict__.items():
+                subname = f"{name}.{key}" if name else key
+                count += count_params(value, subname)
+        elif isinstance(pytree, (list, tuple)):
+            for i, value in enumerate(pytree):
+                subname = f"{name}[{i}]" if name else f"[{i}]"
+                count += count_params(value, subname)
+        elif isinstance(pytree, dict):
+            for key, value in pytree.items():
+                subname = f"{name}.{key}" if name else str(key)
+                count += count_params(value, subname)
+        return count
+
+    count_params(params)
+
+    # Display as table
+    lines = []
+    lines.append("Model Parameter Summary")
+    lines.append("=" * 50)
+    lines.append(f"{'Parameter Name':<30} {'Count':<15}")
+    lines.append("-" * 50)
+
+    for name, count in param_counts.items():
+        lines.append(f"{name:<30} {count:<15,}")
+
+    lines.append("-" * 50)
+    lines.append(f"{'Total Parameters':<30} {total_params:<15,}")
+    lines.append("=" * 50)
+
+    return "\n".join(lines)

noxton-0.1.0/noxton.egg-info/PKG-INFO

@@ -0,0 +1,11 @@
+Metadata-Version: 2.4
+Name: noxton
+Version: 0.1.0
+Summary: Add your description here
+Requires-Python: >=3.13
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: equinox>=0.13.4
+Requires-Dist: jax>=0.9.0.1
+Requires-Dist: statedict2pytree>=2.0.1
+Dynamic: license-file

noxton-0.1.0/noxton.egg-info/SOURCES.txt

@@ -0,0 +1,19 @@
+LICENSE
+README.md
+pyproject.toml
+noxton.egg-info/PKG-INFO
+noxton.egg-info/SOURCES.txt
+noxton.egg-info/dependency_links.txt
+noxton.egg-info/requires.txt
+noxton.egg-info/top_level.txt
+noxton/functions/__init__.py
+noxton/functions/activation.py
+noxton/functions/attention.py
+noxton/functions/embedding.py
+noxton/functions/masking.py
+noxton/functions/normalization.py
+noxton/functions/regularization.py
+noxton/nn/__init__.py
+noxton/nn/abstract.py
+noxton/utils/__init__.py
+noxton/utils/utils.py

noxton-0.1.0/noxton.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

noxton-0.1.0/noxton.egg-info/requires.txt

@@ -0,0 +1,3 @@
+equinox>=0.13.4
+jax>=0.9.0.1
+statedict2pytree>=2.0.1

noxton-0.1.0/noxton.egg-info/top_level.txt

@@ -0,0 +1 @@
+noxton

noxton-0.1.0/pyproject.toml

@@ -0,0 +1,11 @@
+[project]
+name = "noxton"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.13"
+dependencies = [
+    "equinox>=0.13.4",
+    "jax>=0.9.0.1",
+    "statedict2pytree>=2.0.1",
+]

noxton-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+