mlxops-utils 0.1.0__tar.gz

mlxops_utils-0.1.0/.gitignore

@@ -0,0 +1 @@
+.DS_Store

mlxops_utils-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 youyinnn
+
+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.

mlxops_utils-0.1.0/PKG-INFO

@@ -0,0 +1,51 @@
+Metadata-Version: 2.4
+Name: mlxops-utils
+Version: 0.1.0
+Summary: Utilities for MLX operations
+Project-URL: Homepage, https://github.com/youyinnn/mlxops-utils
+Author-email: youyinnn <youyinnn@gmail.com>
+License: MIT License
+        
+        Copyright (c) 2026 youyinnn
+        
+        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.
+License-File: LICENSE
+Requires-Python: >=3.11.4
+Requires-Dist: torch
+Provides-Extra: dev
+Requires-Dist: pytest; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# mlxops-utils
+
+Utilities for MLX operations.
+
+## Installation
+
+```bash
+pip install mlxops-utils
+```
+
+## Requirements
+
+- Python >= 3.11.4
+
+## License
+
+MIT

mlxops_utils-0.1.0/README.md

@@ -0,0 +1,17 @@
+# mlxops-utils
+
+Utilities for MLX operations.
+
+## Installation
+
+```bash
+pip install mlxops-utils
+```
+
+## Requirements
+
+- Python >= 3.11.4
+
+## License
+
+MIT

mlxops_utils-0.1.0/pyproject.toml

@@ -0,0 +1,26 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "mlxops-utils"
+version = "0.1.0"
+description = "Utilities for MLX operations"
+readme = "README.md"
+license = { file = "LICENSE" }
+authors = [
+    { name = "youyinnn", email = "youyinnn@gmail.com" },
+]
+requires-python = ">=3.11.4"
+dependencies = [
+    "torch",
+]
+
+[project.optional-dependencies]
+dev = ["pytest"]
+
+[project.urls]
+Homepage = "https://github.com/youyinnn/mlxops-utils"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/mlxops_utils"]

mlxops_utils-0.1.0/src/mlxops_utils/__init__.py

@@ -0,0 +1,3 @@
+from .data_utils import min_max_normalize, normalize, denormalize, compute_mean_std, IMAGENET_MEAN, IMAGENET_STD
+
+__all__ = ["min_max_normalize", "normalize", "denormalize", "compute_mean_std", "IMAGENET_MEAN", "IMAGENET_STD"]

mlxops_utils-0.1.0/src/mlxops_utils/data_utils.py

@@ -0,0 +1,68 @@
+import torch
+
+
+def min_max_normalize(u: torch.Tensor) -> torch.Tensor:
+    u = u.clone()
+    umin = u.min(dim=-1, keepdim=True).values.min(dim=-2, keepdim=True).values
+    u = u - umin
+    umax = u.max(dim=-1, keepdim=True).values.max(dim=-2, keepdim=True).values
+    u = u / umax
+    if torch.isnan(u).all():
+        u = torch.ones_like(u)
+    return torch.nan_to_num(u)
+
+
+IMAGENET_MEAN = [0.485, 0.456, 0.406]
+IMAGENET_STD = [0.229, 0.224, 0.225]
+
+
+def _build_mean_std(tensors: torch.Tensor, mean: list[float], std: list[float]):
+    c = tensors.shape[1]
+    if len(mean) != c or len(std) != c:
+        raise ValueError(f"Expected mean and std of length {c}, got {len(mean)} and {len(std)}")
+    mean_t = torch.tensor(mean, dtype=tensors.dtype, device=tensors.device).view(1, c, 1, 1)
+    std_t = torch.tensor(std, dtype=tensors.dtype, device=tensors.device).view(1, c, 1, 1)
+    return mean_t, std_t
+
+
+def compute_mean_std(dataloader) -> tuple[list[float], list[float]]:
+    """Compute per-channel mean and std over a dataset.
+
+    Args:
+        dataloader: yields tensors of shape (B, C, H, W), or (tensor, ...) tuples.
+
+    Returns:
+        (mean, std): each a list of floats with length C, ready for use with normalize().
+    """
+    channel_sum = None
+    channel_sq_sum = None
+    pixel_count = 0
+
+    for batch in dataloader:
+        if isinstance(batch, (list, tuple)):
+            batch = batch[0]
+        b, c, h, w = batch.shape
+        if channel_sum is None:
+            channel_sum = torch.zeros(c, dtype=torch.float64)
+            channel_sq_sum = torch.zeros(c, dtype=torch.float64)
+        batch = batch.to(torch.float64)
+        channel_sum += batch.sum(dim=[0, 2, 3])
+        channel_sq_sum += (batch ** 2).sum(dim=[0, 2, 3])
+        pixel_count += b * h * w
+
+    if pixel_count == 0:
+        raise ValueError("dataloader is empty")
+
+    mean = channel_sum / pixel_count
+    std = (channel_sq_sum / pixel_count - mean ** 2).sqrt()
+    return mean.tolist(), std.tolist()
+
+
+def normalize(tensors: torch.Tensor, mean: list[float] = IMAGENET_MEAN, std: list[float] = IMAGENET_STD) -> torch.Tensor:
+    mean_t, std_t = _build_mean_std(tensors, mean, std)
+    return (tensors - mean_t) / std_t
+
+
+def denormalize(tensors: torch.Tensor, mean: list[float] = IMAGENET_MEAN, std: list[float] = IMAGENET_STD) -> torch.Tensor:
+    mean_t, std_t = _build_mean_std(tensors, mean, std)
+    return tensors * std_t + mean_t

mlxops_utils-0.1.0/tests/test_data_utils.py

@@ -0,0 +1,143 @@
+import pytest
+import torch
+from mlxops_utils.data_utils import min_max_normalize, normalize, denormalize, compute_mean_std
+
+
+class TestMinMaxNormMatrix:
+    def test_basic_normalization(self):
+        u = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+        result = min_max_normalize(u)
+        assert result.min().item() == pytest.approx(0.0)
+        assert result.max().item() == pytest.approx(1.0)
+
+    def test_output_range(self):
+        u = torch.rand(5, 5) * 100
+        result = min_max_normalize(u)
+        assert result.min().item() >= 0.0
+        assert result.max().item() <= 1.0
+
+    def test_does_not_modify_input(self):
+        u = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+        original = u.clone()
+        min_max_normalize(u)
+        assert torch.equal(u, original)
+
+    def test_all_same_values_returns_ones(self):
+        # all same -> nan -> torch.isnan(u).all() -> ones_like
+        u = torch.full((3, 3), 5.0)
+        result = min_max_normalize(u)
+        assert torch.equal(result, torch.ones(3, 3))
+
+    def test_4d_batch_tensor(self):
+        # (B, C, H, W): each H×W map normalized independently
+        u = torch.rand(2, 3, 4, 4) * 100
+        result = min_max_normalize(u)
+        assert result.min().item() >= 0.0
+        assert result.max().item() <= 1.0
+        # each (b, c) slice should be in [0, 1]
+        for b in range(2):
+            for c in range(3):
+                s = result[b, c]
+                assert s.min().item() == pytest.approx(0.0)
+                assert s.max().item() == pytest.approx(1.0)
+
+    def test_negative_values(self):
+        u = torch.tensor([[-3.0, -1.0], [0.0, 2.0]])
+        result = min_max_normalize(u)
+        assert result.min().item() == pytest.approx(0.0)
+        assert result.max().item() == pytest.approx(1.0)
+
+    def test_returns_tensor(self):
+        u = torch.rand(3, 3)
+        result = min_max_normalize(u)
+        assert isinstance(result, torch.Tensor)
+
+
+class TestNormalizeDenormalize:
+    def test_norm_output_shape(self):
+        u = torch.rand(2, 3, 64, 64)
+        assert normalize(u).shape == u.shape
+
+    def test_denorm_output_shape(self):
+        u = torch.rand(2, 3, 64, 64)
+        assert denormalize(u).shape == u.shape
+
+    def test_norm_denorm_roundtrip(self):
+        u = torch.rand(2, 3, 64, 64)
+        result = denormalize(normalize(u))
+        assert torch.allclose(result, u, atol=1e-6)
+
+    def test_norm_applies_imagenet_stats(self):
+        # single pixel all-ones image: norm should give (1 - mean) / std per channel
+        u = torch.ones(1, 3, 1, 1)
+        result = normalize(u)
+        expected_r = (1.0 - 0.485) / 0.229
+        assert result[0, 0, 0, 0].item() == pytest.approx(expected_r, rel=1e-5)
+
+    def test_invalid_mean_std_length_raises(self):
+        u = torch.rand(2, 3, 64, 64)
+        with pytest.raises(ValueError):
+            normalize(u, mean=[0.5, 0.5], std=[0.5, 0.5, 0.5])
+
+    def test_custom_mean_std(self):
+        mean = [0.5, 0.5, 0.5]
+        std = [0.5, 0.5, 0.5]
+        u = torch.ones(1, 3, 4, 4)
+        result = normalize(u, mean=mean, std=std)
+        assert result[0, 0, 0, 0].item() == pytest.approx((1.0 - 0.5) / 0.5)
+
+    def test_custom_mean_std_roundtrip(self):
+        mean = [0.3, 0.4, 0.5]
+        std = [0.1, 0.2, 0.3]
+        u = torch.rand(2, 3, 32, 32)
+        assert torch.allclose(denormalize(normalize(u, mean, std), mean, std), u, atol=1e-6)
+
+    def test_single_channel(self):
+        u = torch.rand(2, 1, 64, 64)
+        mean = [0.5]
+        std = [0.25]
+        result = normalize(u, mean=mean, std=std)
+        assert result.shape == u.shape
+
+    def test_does_not_modify_input(self):
+        u = torch.rand(2, 3, 32, 32)
+        original = u.clone()
+        normalize(u)
+        assert torch.equal(u, original)
+
+
+class TestComputeMeanStd:
+    def _make_loader(self, batches):
+        return batches
+
+    def test_known_values(self):
+        # constant tensor: mean = value, std = 0
+        batch = torch.full((4, 3, 8, 8), 0.5)
+        mean, std = compute_mean_std([batch])
+        assert mean == pytest.approx([0.5, 0.5, 0.5], abs=1e-6)
+        assert std == pytest.approx([0.0, 0.0, 0.0], abs=1e-6)
+
+    def test_output_length_matches_channels(self):
+        batch = torch.rand(2, 4, 16, 16)
+        mean, std = compute_mean_std([batch])
+        assert len(mean) == 4
+        assert len(std) == 4
+
+    def test_tuple_batch(self):
+        # dataloader that yields (image, label) tuples
+        batch = (torch.rand(2, 3, 8, 8), torch.zeros(2))
+        mean, std = compute_mean_std([batch])
+        assert len(mean) == 3
+        assert len(std) == 3
+
+    def test_multiple_batches(self):
+        # two batches with same constant: result should match single batch
+        b1 = torch.full((2, 3, 4, 4), 0.3)
+        b2 = torch.full((2, 3, 4, 4), 0.3)
+        mean, std = compute_mean_std([b1, b2])
+        assert mean == pytest.approx([0.3, 0.3, 0.3], abs=1e-6)
+        assert std == pytest.approx([0.0, 0.0, 0.0], abs=1e-6)
+
+    def test_empty_dataloader_raises(self):
+        with pytest.raises(ValueError, match="empty"):
+            compute_mean_std([])