xax 0.0.1__py3-none-any.whl → 0.0.5__py3-none-any.whl

xax/utils/text.py

@@ -0,0 +1,350 @@
+"""Defines helper functions for displaying text in the terminal."""
+
+import datetime
+import itertools
+import re
+import sys
+from typing import Literal
+
+RESET_SEQ = "\033[0m"
+REG_COLOR_SEQ = "\033[%dm"
+BOLD_COLOR_SEQ = "\033[1;%dm"
+BOLD_SEQ = "\033[1m"
+
+Color = Literal[
+    "black",
+    "red",
+    "green",
+    "yellow",
+    "blue",
+    "magenta",
+    "cyan",
+    "white",
+    "grey",
+    "light-red",
+    "light-green",
+    "light-yellow",
+    "light-blue",
+    "light-magenta",
+    "light-cyan",
+]
+
+COLOR_INDEX: dict[Color, int] = {
+    "black": 30,
+    "red": 31,
+    "green": 32,
+    "yellow": 33,
+    "blue": 34,
+    "magenta": 35,
+    "cyan": 36,
+    "white": 37,
+    "grey": 90,
+    "light-red": 91,
+    "light-green": 92,
+    "light-yellow": 93,
+    "light-blue": 94,
+    "light-magenta": 95,
+    "light-cyan": 96,
+}
+
+
+def color_parts(color: Color, bold: bool = False) -> tuple[str, str]:
+    if bold:
+        return BOLD_COLOR_SEQ % COLOR_INDEX[color], RESET_SEQ
+    return REG_COLOR_SEQ % COLOR_INDEX[color], RESET_SEQ
+
+
+def uncolored(s: str) -> str:
+    return re.sub(r"\033\[[\d;]+m", "", s)
+
+
+def colored(s: str, color: Color | None = None, bold: bool = False) -> str:
+    if color is None:
+        return s
+    start, end = color_parts(color, bold=bold)
+    return start + s + end
+
+
+def wrapped(
+    s: str,
+    length: int | None = None,
+    space: str = " ",
+    spaces: str | re.Pattern = r" ",
+    newlines: str | re.Pattern = r"[\n\r]",
+    too_long_suffix: str = "...",
+) -> list[str]:
+    strings = []
+    lines = re.split(newlines, s.strip(), flags=re.MULTILINE | re.UNICODE)
+    for line in lines:
+        cur_string = []
+        cur_length = 0
+        for part in re.split(spaces, line.strip(), flags=re.MULTILINE | re.UNICODE):
+            if length is None:
+                cur_string.append(part)
+                cur_length += len(space) + len(part)
+            else:
+                if len(part) > length:
+                    part = part[: length - len(too_long_suffix)] + too_long_suffix
+                if cur_length + len(part) > length:
+                    strings.append(space.join(cur_string))
+                    cur_string = [part]
+                    cur_length = len(part)
+                else:
+                    cur_string.append(part)
+                    cur_length += len(space) + len(part)
+        if cur_length > 0:
+            strings.append(space.join(cur_string))
+    return strings
+
+
+def outlined(
+    s: str,
+    inner: Color | None = None,
+    side: Color | None = None,
+    bold: bool = False,
+    max_length: int | None = None,
+    space: str = " ",
+    spaces: str | re.Pattern = r" ",
+    newlines: str | re.Pattern = r"[\n\r]",
+) -> str:
+    strs = wrapped(uncolored(s), max_length, space, spaces, newlines)
+    max_len = max(len(s) for s in strs)
+    strs = [f"{s}{' ' * (max_len - len(s))}" for s in strs]
+    strs = [colored(s, inner, bold=bold) for s in strs]
+    strs_with_sides = [f"{colored('│', side)} {s} {colored('│', side)}" for s in strs]
+    top = colored("┌─" + "─" * max_len + "─┐", side)
+    bottom = colored("└─" + "─" * max_len + "─┘", side)
+    return "\n".join([top] + strs_with_sides + [bottom])
+
+
+def show_info(s: str, important: bool = False) -> None:
+    if important:
+        s = outlined(s, inner="light-cyan", side="cyan", bold=True)
+    else:
+        s = colored(s, "light-cyan", bold=False)
+    sys.stdout.write(s)
+    sys.stdout.write("\n")
+    sys.stdout.flush()
+
+
+def show_error(s: str, important: bool = False) -> None:
+    if important:
+        s = outlined(s, inner="light-red", side="red", bold=True)
+    else:
+        s = colored(s, "light-red", bold=False)
+    sys.stdout.write(s)
+    sys.stdout.write("\n")
+    sys.stdout.flush()
+
+
+def show_warning(s: str, important: bool = False) -> None:
+    if important:
+        s = outlined(s, inner="light-yellow", side="yellow", bold=True)
+    else:
+        s = colored(s, "light-yellow", bold=False)
+    sys.stdout.write(s)
+    sys.stdout.write("\n")
+    sys.stdout.flush()
+
+
+class TextBlock:
+    def __init__(
+        self,
+        text: str,
+        color: Color | None = None,
+        bold: bool = False,
+        width: int | None = None,
+        space: str = " ",
+        spaces: str | re.Pattern = r" ",
+        newlines: str | re.Pattern = r"[\n\r]",
+        too_long_suffix: str = "...",
+        no_sep: bool = False,
+        center: bool = False,
+    ) -> None:
+        super().__init__()
+
+        self.width = width
+        self.lines = wrapped(uncolored(text), width, space, spaces, newlines, too_long_suffix)
+        self.color = color
+        self.bold = bold
+        self.no_sep = no_sep
+        self.center = center
+
+
+def render_text_blocks(
+    blocks: list[list[TextBlock]],
+    newline: str = "\n",
+    align_all_blocks: bool = False,
+    padding: int = 0,
+) -> str:
+    """Renders a collection of blocks into a single string.
+
+    Args:
+        blocks: The blocks to render.
+        newline: The string to use as a newline separator.
+        align_all_blocks: If set, aligns the widths for all blocks.
+        padding: The amount of padding to add to each block.
+
+    Returns:
+        The rendered blocks.
+    """
+    if align_all_blocks:
+        if any(len(row) != len(blocks[0]) for row in blocks):
+            raise ValueError("All rows must have the same number of blocks in order to align them")
+        widths = [[max(len(line) for line in i.lines) if i.width is None else i.width for i in r] for r in blocks]
+        row_widths = [max(i) for i in zip(*widths)]
+        for row in blocks:
+            for i, block in enumerate(row):
+                block.width = row_widths[i]
+
+    def get_widths(row: list[TextBlock], n: int = 0) -> list[int]:
+        return [
+            (max(len(line) for line in block.lines) if block.width is None else block.width) + n + padding
+            for block in row
+        ]
+
+    def get_acc_widths(row: list[TextBlock], n: int = 0) -> list[int]:
+        return list(itertools.accumulate(get_widths(row, n)))
+
+    def get_height(row: list[TextBlock]) -> int:
+        return max(len(block.lines) for block in row)
+
+    def pad(s: str, width: int, center: bool) -> str:
+        swidth = len(s)
+        if center:
+            lpad, rpad = (width - swidth) // 2, (width - swidth + 1) // 2
+        else:
+            lpad, rpad = 0, width - swidth
+        return " " * lpad + s + " " * rpad
+
+    lines = []
+    prev_row: list[TextBlock] | None = None
+    for row in blocks:
+        if prev_row is None:
+            lines += ["┌─" + "─┬─".join(["─" * width for width in get_widths(row)]) + "─┐"]
+        elif not all(block.no_sep for block in row):
+            ins, outs = get_acc_widths(prev_row, 3), get_acc_widths(row, 3)
+            segs = sorted([(i, False) for i in ins] + [(i, True) for i in outs])
+            line = ["├"]
+
+            c = 1
+            for i, (s, is_out) in enumerate(segs):
+                if i > 0 and segs[i - 1][0] == s:
+                    continue
+                is_in_out = i < len(segs) - 1 and segs[i + 1][0] == s
+                is_last = i == len(segs) - 2 if is_in_out else i == len(segs) - 1
+
+                line += "─" * (s - c)
+                if is_last:
+                    if is_in_out:
+                        line += "┤"
+                    elif is_out:
+                        line += "┐"
+                    else:
+                        line += "┘"
+                else:  # noqa: PLR5501
+                    if is_in_out:
+                        line += "┼"
+                    elif is_out:
+                        line += "┬"
+                    else:
+                        line += "┴"
+                c = s + 1
+
+            lines += ["".join(line)]
+
+        for i in range(get_height(row)):
+            lines += [
+                "│ "
+                + " │ ".join(
+                    [
+                        (
+                            " " * width
+                            if i >= len(block.lines)
+                            else colored(pad(block.lines[i], width, block.center), block.color, bold=block.bold)
+                        )
+                        for block, width in zip(row, get_widths(row))
+                    ]
+                )
+                + " │"
+            ]
+
+        prev_row = row
+    if prev_row is not None:
+        lines += ["└─" + "─┴─".join(["─" * width for width in get_widths(prev_row)]) + "─┘"]
+
+    return newline.join(lines)
+
+
+def format_timedelta(timedelta: datetime.timedelta, short: bool = False) -> str:
+    """Formats a delta time to human-readable format.
+
+    Args:
+        timedelta: The delta to format
+        short: If set, uses a shorter format
+
+    Returns:
+        The human-readable time delta
+    """
+    parts = []
+    if timedelta.days > 0:
+        if short:
+            parts += [f"{timedelta.days}d"]
+        else:
+            parts += [f"{timedelta.days} day" if timedelta.days == 1 else f"{timedelta.days} days"]
+
+    seconds = timedelta.seconds
+
+    if seconds > 60 * 60:
+        hours, seconds = seconds // (60 * 60), seconds % (60 * 60)
+        if short:
+            parts += [f"{hours}h"]
+        else:
+            parts += [f"{hours} hour" if hours == 1 else f"{hours} hours"]
+
+    if seconds > 60:
+        minutes, seconds = seconds // 60, seconds % 60
+        if short:
+            parts += [f"{minutes}m"]
+        else:
+            parts += [f"{minutes} minute" if minutes == 1 else f"{minutes} minutes"]
+
+    if short:
+        parts += [f"{seconds}s"]
+    else:
+        parts += [f"{seconds} second" if seconds == 1 else f"{seconds} seconds"]
+
+    return ", ".join(parts)
+
+
+def format_datetime(dt: datetime.datetime) -> str:
+    """Formats a datetime to human-readable format.
+
+    Args:
+        dt: The datetime to format
+
+    Returns:
+        The human-readable datetime
+    """
+    return dt.strftime("%Y-%m-%d %H:%M:%S")
+
+
+def camelcase_to_snakecase(s: str) -> str:
+    return re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", s).lower()
+
+
+def snakecase_to_camelcase(s: str) -> str:
+    return "".join(word.title() for word in s.split("_"))
+
+
+def highlight_exception_message(s: str) -> str:
+    s = re.sub(r"(\w+Error)", r"\033[1;31m\1\033[0m", s)
+    s = re.sub(r"(\w+Exception)", r"\033[1;31m\1\033[0m", s)
+    s = re.sub(r"(\w+Warning)", r"\033[1;33m\1\033[0m", s)
+    s = re.sub(r"\^+", r"\033[1;35m\g<0>\033[0m", s)
+    s = re.sub(r"File \"(.+?)\"", r'File "\033[36m\1\033[0m"', s)
+    return s
+
+
+def is_interactive_session() -> bool:
+    return hasattr(sys, "ps1") or hasattr(sys, "ps2") or sys.stdout.isatty() or sys.stderr.isatty()

xax-0.0.5.dist-info/METADATA

@@ -0,0 +1,40 @@
+Metadata-Version: 2.1
+Name: xax
+Version: 0.0.5
+Summary: The xax project
+Home-page: https://github.com/dpshai/xax
+Author: Benjamin Bolte
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax
+Requires-Dist: jaxtyping
+Requires-Dist: equinox
+Requires-Dist: optax
+Requires-Dist: dpshdl
+Requires-Dist: cloudpickle
+Requires-Dist: pillow
+Requires-Dist: omegaconf
+Requires-Dist: gitpython
+Requires-Dist: tensorboard
+Requires-Dist: psutil
+Requires-Dist: requests
+Provides-Extra: dev
+Requires-Dist: black; extra == "dev"
+Requires-Dist: darglint; extra == "dev"
+Requires-Dist: mypy; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: types-pillow; extra == "dev"
+Requires-Dist: types-psutil; extra == "dev"
+Requires-Dist: types-requests; extra == "dev"
+
+# xax
+
+JAX library for fast experimentation.
+
+## Installation
+
+```bash
+pip install xax
+```

xax-0.0.5.dist-info/RECORD

@@ -0,0 +1,52 @@
+xax/__init__.py,sha256=3OQTnHGYgaux3i9gTYZxfK8F2zS_hK2QqD-G-Z1TfHQ,7623
+xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
+xax/requirements.txt,sha256=DRn2B9d3mAr57-U3IOIrKm2nYz8H3cYgDy6EIC3SsuE,266
+xax/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xax/core/conf.py,sha256=hwgc5sJw0YRSegQLLrmIDtscev-H_a2ST1-V6BJ5aec,5915
+xax/core/state.py,sha256=7lnVSytuhwPfcobPGdjfQ0QxbLgzWQNipKwXchd58QI,2695
+xax/nn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xax/nn/embeddings.py,sha256=bQGxBFxkLwi2MQLkRfGaHPH5P_KKB21HdI7VNWTKIOQ,11847
+xax/nn/functions.py,sha256=CI_OmspaQwN9nl4hwefIU3_I7m6gBZwJ9aGK1JGUgr0,2713
+xax/nn/parallel.py,sha256=fnTiT7MsG7eQrJvqwjIz2Ifo3P27TuxIJzmpGYSa_dQ,4608
+xax/task/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xax/task/base.py,sha256=n82Sw-kMLr-WZzh0c_vAAQ2b-DHRYs0U8biPRonBxKU,7252
+xax/task/logger.py,sha256=MAFIgd6yO0pD3gJHfKTwUDcwaM8DZD3AZtFLvrQtlFo,26740
+xax/task/script.py,sha256=oBGnScYa_X284fCajabPCcbaSEIqR8nO4d40dvMv3NQ,1011
+xax/task/task.py,sha256=X7TV_gt6C4m_-Il22Uyr5iMm-eh15oH5v1dl96sv1go,1295
+xax/task/launchers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xax/task/launchers/base.py,sha256=8LB_r6YISKu1vq1zk3aVYmiedRr9MxE5IMRocs6unFI,731
+xax/task/launchers/cli.py,sha256=cK7Nm-3fO-W2gTxpn3FEThsT2NvneS2w0UjA1Nt-84A,1402
+xax/task/launchers/single_process.py,sha256=IoML-30g5c526yxkpbWSOtG_KpNQMakT7xujzB1gIAo,846
+xax/task/launchers/staged.py,sha256=jYeT9u58CN4ldV-ltJiQXQglEWOnEckHWnHYjfJQaoY,1102
+xax/task/loggers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xax/task/loggers/callback.py,sha256=reaRuJs5iB6WWNgh3_tsuz_QPAlBC-5Ed2wCG_6Wj4M,2075
+xax/task/loggers/json.py,sha256=yXHb1bmfsEnk4p0F1Up1ertWYdcPAFZm25NT8wE3Jb8,4045
+xax/task/loggers/state.py,sha256=qyb-q8MdagN7BX-DhKucwoc45tIZJrPuvVDVoysTKC4,1576
+xax/task/loggers/stdout.py,sha256=nxQXkS9JUR38RKsU9qj0dgePKguK0BFa9nl_BdGO8cE,6758
+xax/task/loggers/tensorboard.py,sha256=DMYRDCQ9c-xHqO4kkZvc1-53PXCf2gX0aRiiAQDtHJ0,7293
+xax/task/mixins/__init__.py,sha256=NkSAjMN5jpXE6LROIwMzX60z7UsTBpGs624_mNUWquo,745
+xax/task/mixins/artifacts.py,sha256=G0984WuXII_R13IlJZn9En7iM83ISXKjeVYvn7j4wBs,3754
+xax/task/mixins/checkpointing.py,sha256=JV91b5xyBUyZIbR3S-5UkBZNoAZYCnWx7Y-ayuU0lHQ,7989
+xax/task/mixins/cpu_stats.py,sha256=Lqskt1t4usE6UslhANjwB0ZKOYmaC4dm9dnVKa6ERdA,8924
+xax/task/mixins/data_loader.py,sha256=BPs0sYdctesnhS9nQ1rvT77MzLXznw5E4tAzWT1PpJY,5998
+xax/task/mixins/gpu_stats.py,sha256=tFTNmtl9iMiLiYJSPg7gHR-ZxOP4P_ynzSmYNIAUoRw,8431
+xax/task/mixins/logger.py,sha256=6XkjP_YUGY2CiDry0kDm1f9jqzJaLa1bPVYYnGjvSBU,2049
+xax/task/mixins/process.py,sha256=HQAvEruvvfcS_IThrM4hKhFHZCAN2kFY_vEaZGLeZS8,1428
+xax/task/mixins/runnable.py,sha256=d5-qyIpmNPtbTzE7qFJGGCPSREEDhX1VApUJPNDWye0,1933
+xax/task/mixins/step_wrapper.py,sha256=Do4eGgZVuqDX9ZGDxQdfn6pRbUnHjQBAkTF0vnNH31E,1472
+xax/task/mixins/train.py,sha256=Xeb0N9j-Znz5QnMDCXDGPqUSKMNLJkd8oF8giN45l2U,20099
+xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xax/utils/experiments.py,sha256=qT3H0fyVH8DN417x7T0Xmz4SKoogW81-EHcZfyktFI8,28300
+xax/utils/jax.py,sha256=VzEVB766UyH3_cgN6UP0FkCsDuGlYg5KJj8YJS4yYUk,439
+xax/utils/logging.py,sha256=ST1hp2C2xntVVJBUHwo3YxPK19fBLNvHU2WGO1xqcXA,6418
+xax/utils/numpy.py,sha256=_jOXVi-d2AtJnRftPkRK5MDMzsU8slgw-Jjv4GRm6ns,1197
+xax/utils/tensorboard.py,sha256=XqxUlryFVsb75jE36uLcuoUhSr3nWg_-dzji2h6U_rI,8245
+xax/utils/text.py,sha256=zo1sAoZe59GkpcpaHBVOQ0OekSMGXvOAyNa3lOJozCY,10628
+xax/utils/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xax/utils/data/collate.py,sha256=Rd9vMomr_S_zCa_Hi4dO-8ntzAfVwndIUtuXFA3iNcc,7066
+xax-0.0.5.dist-info/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.0.5.dist-info/METADATA,sha256=VCiQmbjwZtiuORVyB0dloFTgLWtnK4o3FaolNWvf-A4,937
+xax-0.0.5.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+xax-0.0.5.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.0.5.dist-info/RECORD,,

{xax-0.0.1.dist-info → xax-0.0.5.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.42.0)
+Generator: setuptools (75.6.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

xax-0.0.5.dist-info/top_level.txt

@@ -0,0 +1 @@
+xax

examples/mnist.py

@@ -1,148 +0,0 @@
-"""MNIST example in Jax."""
-
-import array
-import gzip
-import itertools
-import os
-import struct
-import time
-import urllib.request
-from typing import Iterator
-
-import jax.numpy as jnp
-import numpy as np
-import numpy.random as npr
-from jax import grad, jit, random
-from jax.example_libraries import optimizers, stax
-from jax.example_libraries.optimizers import OptimizerState
-from jax.example_libraries.stax import Dense, LogSoftmax, Relu
-from jaxtyping import ArrayLike
-
-_DATA = "/tmp/jax_example_data/"
-
-
-def _download(url: str, filename: str) -> None:
-    if not os.path.exists(_DATA):
-        os.makedirs(_DATA)
-    out_file = os.path.join(_DATA, filename)
-    if not os.path.isfile(out_file):
-        urllib.request.urlretrieve(url, out_file)
-        print(f"downloaded {url} to {_DATA}")
-
-
-def _partial_flatten(x: np.ndarray) -> np.ndarray:
-    return np.reshape(x, (x.shape[0], -1))
-
-
-def _one_hot(x: np.ndarray, k: int, dtype: type = np.float32) -> np.ndarray:
-    return np.array(x[:, None] == np.arange(k), dtype)
-
-
-def mnist_raw() -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
-    base_url = "https://storage.googleapis.com/cvdf-datasets/mnist/"
-
-    def parse_labels(filename: str) -> np.ndarray:
-        with gzip.open(filename, "rb") as fh:
-            _ = struct.unpack(">II", fh.read(8))
-            return np.array(array.array("B", fh.read()), dtype=np.uint8)
-
-    def parse_images(filename: str) -> np.ndarray:
-        with gzip.open(filename, "rb") as fh:
-            _, num_data, rows, cols = struct.unpack(">IIII", fh.read(16))
-            return np.array(array.array("B", fh.read()), dtype=np.uint8).reshape(num_data, rows, cols)
-
-    for filename in [
-        "train-images-idx3-ubyte.gz",
-        "train-labels-idx1-ubyte.gz",
-        "t10k-images-idx3-ubyte.gz",
-        "t10k-labels-idx1-ubyte.gz",
-    ]:
-        _download(base_url + filename, filename)
-
-    train_images = parse_images(os.path.join(_DATA, "train-images-idx3-ubyte.gz"))
-    train_labels = parse_labels(os.path.join(_DATA, "train-labels-idx1-ubyte.gz"))
-    test_images = parse_images(os.path.join(_DATA, "t10k-images-idx3-ubyte.gz"))
-    test_labels = parse_labels(os.path.join(_DATA, "t10k-labels-idx1-ubyte.gz"))
-
-    return train_images, train_labels, test_images, test_labels
-
-
-def mnist(permute_train: bool = False) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
-    train_images, train_labels, test_images, test_labels = mnist_raw()
-
-    train_images = _partial_flatten(train_images) / np.float32(255.0)
-    test_images = _partial_flatten(test_images) / np.float32(255.0)
-    train_labels = _one_hot(train_labels, 10)
-    test_labels = _one_hot(test_labels, 10)
-
-    if permute_train:
-        perm = np.random.RandomState(0).permutation(train_images.shape[0])
-        train_images = train_images[perm]
-        train_labels = train_labels[perm]
-
-    return train_images, train_labels, test_images, test_labels
-
-
-def loss(params: tuple[ArrayLike, ArrayLike], batch: tuple[ArrayLike, ArrayLike]) -> ArrayLike:
-    inputs, targets = batch
-    preds = predict(params, inputs)
-    return -jnp.mean(jnp.sum(preds * targets, axis=1))
-
-
-def accuracy(params: tuple[ArrayLike, ArrayLike], batch: tuple[ArrayLike, ArrayLike]) -> ArrayLike:
-    inputs, targets = batch
-    target_class = jnp.argmax(targets, axis=1)
-    predicted_class = jnp.argmax(predict(params, inputs), axis=1)
-    return jnp.mean(predicted_class == target_class)
-
-
-init_random_params, predict = stax.serial(Dense(1024), Relu, Dense(1024), Relu, Dense(10), LogSoftmax)
-
-if __name__ == "__main__":
-    # python -m examples.mnist
-    rng = random.PRNGKey(0)
-
-    step_size = 0.001
-    num_epochs = 10
-    batch_size = 128
-    momentum_mass = 0.9
-
-    train_images, train_labels, test_images, test_labels = mnist()
-    num_train = train_images.shape[0]
-    num_complete_batches, leftover = divmod(num_train, batch_size)
-    num_batches = num_complete_batches + bool(leftover)
-
-    def data_stream() -> Iterator[tuple[np.ndarray, np.ndarray]]:
-        rng = npr.RandomState(0)
-        while True:
-            perm = rng.permutation(num_train)
-            for i in range(num_batches):
-                batch_idx = perm[i * batch_size : (i + 1) * batch_size]
-                yield train_images[batch_idx], train_labels[batch_idx]
-
-    batches = data_stream()
-
-    opt_init, opt_update, get_params = optimizers.momentum(step_size, mass=momentum_mass)
-
-    @jit
-    def update(i: int, opt_state: OptimizerState, batch: tuple[ArrayLike, ArrayLike]) -> OptimizerState:
-        params = get_params(opt_state)
-        return opt_update(i, grad(loss)(params, batch), opt_state)
-
-    _, init_params = init_random_params(rng, (-1, 28 * 28))
-    opt_state = opt_init(init_params)
-    itercount = itertools.count()
-
-    print("\nStarting training...")
-    for epoch in range(num_epochs):
-        start_time = time.time()
-        for _ in range(num_batches):
-            opt_state = update(next(itercount), opt_state, next(batches))
-        epoch_time = time.time() - start_time
-
-        params = get_params(opt_state)
-        train_acc = accuracy(params, (train_images, train_labels))
-        test_acc = accuracy(params, (test_images, test_labels))
-        print(f"Epoch {epoch} in {epoch_time:0.2f} sec")
-        print(f"Training set accuracy {train_acc}")
-        print(f"Test set accuracy {test_acc}")

xax-0.0.1.dist-info/METADATA

@@ -1,21 +0,0 @@
-Metadata-Version: 2.1
-Name: xax
-Version: 0.0.1
-Summary: The xax project
-Home-page: https://github.com/dpshai/xax
-Author: Benjamin Bolte
-Requires-Python: >=3.11
-Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: jax
-Requires-Dist: jaxtyping
-Provides-Extra: dev
-Requires-Dist: black ; extra == 'dev'
-Requires-Dist: darglint ; extra == 'dev'
-Requires-Dist: mypy ; extra == 'dev'
-Requires-Dist: pytest ; extra == 'dev'
-Requires-Dist: ruff ; extra == 'dev'
-
-# xax
-
-JAX library for fast experimentation.

xax-0.0.1.dist-info/RECORD

@@ -1,9 +0,0 @@
-examples/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-examples/mnist.py,sha256=vIKKlxS153hWM0Z8xM3HzbbF1leWrgRBPzL_fDoTMko,5366
-xax/__init__.py,sha256=sXLh7g3KC4QCFxcZGBTpG2scR7hmmBsMjq6LqRptkRg,22
-xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax-0.0.1.dist-info/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.0.1.dist-info/METADATA,sha256=F9reokda5s_h_HTX5dKfUKRxNKLoAWNs5nnIR-g9kj0,524
-xax-0.0.1.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-xax-0.0.1.dist-info/top_level.txt,sha256=rD77ScL4NoinNhAT1cDqumVhFa2AoqoIrgBKbwSsOzY,13
-xax-0.0.1.dist-info/RECORD,,

xax-0.0.1.dist-info/top_level.txt

@@ -1,2 +0,0 @@
-examples
-xax