waxjax 0.1.0__tar.gz

waxjax-0.1.0/MANIFEST.in

@@ -0,0 +1,2 @@
+include README.md
+global-exclude __pycache__

waxjax-0.1.0/PKG-INFO

@@ -0,0 +1,55 @@
+Metadata-Version: 2.4
+Name: waxjax
+Version: 0.1.0
+Summary: Load HuggingFace SafeTensors models to JAX/Flax
+Home-page: https://github.com/ik12-b/waxjax
+Author: Your Name
+License: Apache-2.0
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: safetensors>=0.4.0
+Requires-Dist: numpy>=1.24
+Requires-Dist: jax>=0.4.0
+Provides-Extra: torch
+Requires-Dist: torch>=2.0; extra == "torch"
+Provides-Extra: flax
+Requires-Dist: flax>=0.7; extra == "flax"
+Provides-Extra: all
+Requires-Dist: torch>=2.0; extra == "all"
+Requires-Dist: flax>=0.7; extra == "all"
+Requires-Dist: ml_dtypes; extra == "all"
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: black; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Dynamic: author
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# waxjax
+
+Lightweight helper to load HuggingFace SafeTensors checkpoints into JAX/Flax.
+
+Install (latest from source):
+
+```bash
+pip install git+https://github.com/ik12-b/waxjax.git
+```
+
+Build and publish (recommended to test to TestPyPI first): see `RELEASE.md` or run locally:
+
+```bash
+python -m pip install --upgrade build twine
+python -m build
+twine upload --repository testpypi dist/*
+```

waxjax-0.1.0/README.md

@@ -0,0 +1,17 @@
+# waxjax
+
+Lightweight helper to load HuggingFace SafeTensors checkpoints into JAX/Flax.
+
+Install (latest from source):
+
+```bash
+pip install git+https://github.com/ik12-b/waxjax.git
+```
+
+Build and publish (recommended to test to TestPyPI first): see `RELEASE.md` or run locally:
+
+```bash
+python -m pip install --upgrade build twine
+python -m build
+twine upload --repository testpypi dist/*
+```

waxjax-0.1.0/pyproject.toml

@@ -0,0 +1,3 @@
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"

waxjax-0.1.0/safejax/__init__.py

@@ -0,0 +1,9 @@
+from .core.api import load, save, inspect, verify
+from .core.rules import MappingRule
+from .architectures._registry import register_architecture, ArchitectureConfig
+
+__version__ = "0.1.0"
+__all__ = [
+    "load", "save", "inspect", "verify",
+    "MappingRule", "register_architecture", "ArchitectureConfig",
+]

waxjax-0.1.0/setup.cfg

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

waxjax-0.1.0/setup.py

@@ -0,0 +1,34 @@
+from pathlib import Path
+from setuptools import setup, find_packages
+
+here = Path(__file__).parent
+long_description = (here / "README.md").read_text(encoding="utf-8") if (here / "README.md").exists() else ""
+
+setup(
+    name="waxjax",
+    version="0.1.0",
+    description="Load HuggingFace SafeTensors models to JAX/Flax",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    author="Your Name",
+    license="Apache-2.0",
+    url="https://github.com/ik12-b/waxjax",
+    packages=find_packages(),
+    python_requires=">=3.9",
+    install_requires=[
+        "safetensors>=0.4.0",
+        "numpy>=1.24",
+        "jax>=0.4.0",
+    ],
+    extras_require={
+        "torch": ["torch>=2.0"],
+        "flax": ["flax>=0.7"],
+        "all": ["torch>=2.0", "flax>=0.7", "ml_dtypes"],
+        "dev": ["pytest", "black", "ruff"],
+    },
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "License :: OSI Approved :: Apache Software License",
+        "Operating System :: OS Independent",
+    ],
+)

waxjax-0.1.0/tests/test_api.py

@@ -0,0 +1,72 @@
+# tests/test_api.py
+"""Test public API: inspect, verify, unsupported model."""
+
+import pytest
+import json
+import waxjax
+
+
+class TestInspect:
+
+    def test_inspect_runs_without_error(self, qwen2_dir, capsys):
+        waxjax.inspect(qwen2_dir)
+        out = capsys.readouterr().out
+        assert "qwen2" in out
+        assert "✓" in out
+
+    def test_inspect_shows_unsupported(self, tmp_path, capsys):
+        (tmp_path / "config.json").write_text(
+            json.dumps({"model_type": "unknown_model_xyz"})
+        )
+        # Buat dummy safetensors
+        import torch
+        from safetensors.torch import save_file
+        save_file({"w": torch.randn(2,2)}, str(tmp_path / "model.safetensors"))
+
+        waxjax.inspect(str(tmp_path))
+        out = capsys.readouterr().out
+        assert "✗" in out or "belum didukung" in out
+
+
+class TestUnsupportedModel:
+
+    def test_load_unknown_raises_valueerror(self, tmp_path):
+        (tmp_path / "config.json").write_text(
+            json.dumps({"model_type": "unknown_xyz"})
+        )
+        import torch
+        from safetensors.torch import save_file
+        save_file({"w": torch.randn(2,2)}, str(tmp_path / "model.safetensors"))
+
+        with pytest.raises(ValueError, match="belum didukung"):
+            waxjax.load(str(tmp_path))
+
+    def test_load_with_custom_rules_bypasses_registry(self, tmp_path):
+        """Custom rules bisa dipakai untuk model yang belum terdaftar."""
+        from waxjax.core.rules import MappingRule
+        import torch
+        from safetensors.torch import save_file
+
+        (tmp_path / "config.json").write_text(
+            json.dumps({"model_type": "custom_unknown"})
+        )
+        save_file({"fc.weight": torch.randn(4, 8)},
+                  str(tmp_path / "model.safetensors"))
+
+        rules = [
+            MappingRule("fc.weight",
+                        rename=lambda k: k.replace(".weight", ".kernel"),
+                        transform="transpose_2d")
+        ]
+        params = waxjax.load(str(tmp_path), rules=rules)
+        assert "fc" in params
+        assert "kernel" in params["fc"]
+
+
+class TestVerify:
+
+    def test_verify_passes_for_correct_conversion(self, qwen2_dir):
+        params = waxjax.load(qwen2_dir)
+        report = waxjax.verify(params, qwen2_dir, sample_n=5)
+        assert report["passed"] == report["total"]
+        assert report["max_diff"] < 1e-5

waxjax-0.1.0/tests/test_loader.py

@@ -0,0 +1,65 @@
+# tests/test_loader.py
+"""Test loader: baca safetensors, handle bfloat16, sharding."""
+
+import pytest
+import numpy as np
+from pathlib import Path
+from waxjax.core.loader import load_safetensors
+
+
+class TestLoader:
+
+    def test_load_basic(self, qwen2_dir):
+        weights = load_safetensors(qwen2_dir)
+        assert isinstance(weights, dict)
+        assert len(weights) > 0
+
+    def test_all_float32(self, qwen2_dir):
+        """Semua tensor harus float32 setelah load."""
+        weights = load_safetensors(qwen2_dir)
+        for key, arr in weights.items():
+            assert arr.dtype == np.float32, f"{key} bukan float32: {arr.dtype}"
+
+    def test_expected_keys_present_qwen2(self, qwen2_dir):
+        weights = load_safetensors(qwen2_dir)
+        assert "model.embed_tokens.weight" in weights
+        assert "model.layers.0.self_attn.q_proj.weight" in weights
+        assert "model.layers.0.self_attn.q_proj.bias" in weights
+        assert "model.norm.weight" in weights
+
+    def test_expected_keys_present_nllb(self, nllb_dir):
+        weights = load_safetensors(nllb_dir)
+        assert "model.shared.weight" in weights
+        assert "model.encoder.layers.0.self_attn.q_proj.weight" in weights
+        assert "model.decoder.layers.0.encoder_attn.q_proj.weight" in weights
+        assert "model.encoder.embed_positions.weight" in weights
+
+    def test_missing_dir_raises(self, tmp_path):
+        with pytest.raises(FileNotFoundError):
+            load_safetensors(str(tmp_path / "nonexistent"))
+
+    def test_bfloat16_handling(self, tmp_path):
+        """Tensor bfloat16 harus di-cast ke float32 tanpa error."""
+        import torch
+        from safetensors.torch import save_file
+
+        bf16_tensor = torch.randn(4, 8).to(torch.bfloat16)
+        save_file({"layer.weight": bf16_tensor}, str(tmp_path / "model.safetensors"))
+
+        weights = load_safetensors(str(tmp_path))
+        assert weights["layer.weight"].dtype == np.float32
+
+    def test_sharded_loading(self, tmp_path):
+        """Dua file safetensors di-merge dengan benar."""
+        import torch
+        from safetensors.torch import save_file
+
+        save_file({"shard1.weight": torch.randn(4, 4)},
+                  str(tmp_path / "model-00001-of-00002.safetensors"))
+        save_file({"shard2.weight": torch.randn(4, 4)},
+                  str(tmp_path / "model-00002-of-00002.safetensors"))
+
+        weights = load_safetensors(str(tmp_path))
+        assert "shard1.weight" in weights
+        assert "shard2.weight" in weights
+        assert len(weights) == 2

waxjax-0.1.0/tests/test_mapper.py

@@ -0,0 +1,88 @@
+# tests/test_mapper.py
+"""Test mapper: apply rules, skip, tied weights."""
+
+import pytest
+import numpy as np
+from waxjax.core.mapper import apply_rules
+from waxjax.core.rules import MappingRule
+
+
+def _linear_rules():
+    return [
+        MappingRule("**.weight", rename=lambda k: k.replace(".weight", ".kernel"),
+                    transform="transpose_2d"),
+        MappingRule("**.bias",   rename=lambda k: k, transform="no_op"),
+    ]
+
+
+class TestMapper:
+
+    def test_basic_rename_and_transpose(self):
+        weights = {"fc.weight": np.ones((4, 8))}
+        result  = apply_rules(weights, _linear_rules())
+        assert ("fc", "kernel") in result
+        assert result[("fc", "kernel")].shape == (8, 4)
+
+    def test_bias_not_transposed(self):
+        weights = {"fc.bias": np.ones(8)}
+        result  = apply_rules(weights, _linear_rules())
+        assert result[("fc", "bias")].shape == (8,)
+
+    def test_skip_sentinel(self):
+        """Key dengan rename → '__skip__' tidak masuk ke output."""
+        rules = [
+            MappingRule(
+                "tied.weight",
+                rename=lambda k: "__skip__",
+                transform="no_op",
+            )
+        ]
+        weights = {"tied.weight": np.ones((4, 4))}
+        result  = apply_rules(weights, rules)
+        assert ("tied", "weight") not in result
+        assert len(result) == 0
+
+    def test_priority_order(self):
+        """Rule lebih spesifik menang atas rule lebih umum."""
+        specific = MappingRule(
+            "model.norm.weight",
+            rename=lambda k: k.replace(".weight", ".scale"),
+            transform="no_op",
+        )
+        generic = MappingRule(
+            "**.weight",
+            rename=lambda k: k.replace(".weight", ".kernel"),
+            transform="transpose_2d",
+        )
+        weights = {"model.norm.weight": np.ones(8)}
+        result  = apply_rules(weights, [generic, specific])
+        # Harus pakai rule spesifik → .scale, bukan .kernel
+        assert ("model", "norm", "scale") in result
+        assert ("model", "norm", "kernel") not in result
+
+    def test_tied_weights_auto_created(self):
+        """Tied weight dibuat otomatis jika source ada tapi target tidak."""
+        rules = [
+            MappingRule(
+                "model.shared.weight",
+                rename=lambda k: k.replace(".weight", ".embedding"),
+                transform="no_op",
+            )
+        ]
+        weights = {"model.shared.weight": np.ones((512, 64))}
+        result  = apply_rules(
+            weights, rules,
+            tied_weights={"model.shared.weight": "lm_head.kernel"}
+        )
+        # embedding harus ada
+        assert ("model", "shared", "embedding") in result
+        # lm_head.kernel harus dibuat otomatis dengan transpose
+        assert ("lm_head", "kernel") in result
+        assert result[("lm_head", "kernel")].shape == (64, 512)
+
+    def test_unmatched_keys_warn(self):
+        """Key tanpa rule matching menghasilkan warning, tidak error."""
+        weights = {"unknown.mystery": np.ones((4, 4))}
+        with pytest.warns(UserWarning, match="tidak match"):
+            result = apply_rules(weights, _linear_rules())
+        assert len(result) == 0

waxjax-0.1.0/tests/test_nllb.py

@@ -0,0 +1,93 @@
+# tests/test_nllb.py
+"""Integration test end-to-end untuk arsitektur NLLB/M2M-100."""
+
+import pytest
+import numpy as np
+from flax.traverse_util import flatten_dict
+import waxjax
+
+
+class TestNLLBLoad:
+
+    def test_load_success(self, nllb_dir):
+        params = waxjax.load(nllb_dir)
+        assert isinstance(params, dict)
+
+    def test_shared_embedding_shape(self, nllb_dir):
+        """model.shared.embedding shape (V, D) tidak di-transpose."""
+        params = waxjax.load(nllb_dir)
+        embed  = np.array(params["model"]["shared"]["embedding"])
+        assert embed.shape == (512, 64)
+
+    def test_tied_embed_tokens_skipped(self, nllb_dir):
+        """encoder/decoder embed_tokens tidak masuk params (tied ke shared)."""
+        params = waxjax.load(nllb_dir)
+        flat   = flatten_dict(params)
+        assert ("model","encoder","embed_tokens","embedding") not in flat
+        assert ("model","decoder","embed_tokens","embedding") not in flat
+
+    def test_positional_embedding_present(self, nllb_dir):
+        params = waxjax.load(nllb_dir)
+        flat   = flatten_dict(params)
+        assert ("model","encoder","embed_positions","embedding") in flat
+        assert ("model","decoder","embed_positions","embedding") in flat
+        # Shape (MAXPOS+2, D) = (132, 64)
+        pos = np.array(flat[("model","encoder","embed_positions","embedding")])
+        assert pos.shape == (132, 64)
+
+    def test_layernorm_has_scale_and_bias(self, nllb_dir):
+        """NLLB pakai LayerNorm → punya scale DAN bias (beda dari Qwen2 RMSNorm)."""
+        params = waxjax.load(nllb_dir)
+        flat   = flatten_dict(params)
+        assert ("model","encoder","layers","0","self_attn_layer_norm","scale") in flat
+        assert ("model","encoder","layers","0","self_attn_layer_norm","bias")  in flat
+
+    def test_fc1_fc2_kernel_transposed(self, nllb_dir):
+        """fc1/fc2 harus di-transpose: PT (FFN,D) → Flax (D,FFN)."""
+        params = waxjax.load(nllb_dir)
+        flat   = flatten_dict(params)
+        fc1 = np.array(flat[("model","encoder","layers","0","fc1","kernel")])
+        fc2 = np.array(flat[("model","encoder","layers","0","fc2","kernel")])
+        # D=64, FFN=128
+        assert fc1.shape == (64, 128)   # transposed dari (128,64)
+        assert fc2.shape == (128, 64)   # transposed dari (64,128)
+
+    def test_out_proj_not_o_proj(self, nllb_dir):
+        """NLLB pakai out_proj bukan o_proj seperti Qwen2."""
+        params = waxjax.load(nllb_dir)
+        flat   = flatten_dict(params)
+        assert ("model","encoder","layers","0","self_attn","out_proj","kernel") in flat
+        assert ("model","encoder","layers","0","self_attn","o_proj","kernel") not in flat
+
+    def test_cross_attention_present(self, nllb_dir):
+        """Decoder harus punya encoder_attn (cross-attention)."""
+        params = waxjax.load(nllb_dir)
+        flat   = flatten_dict(params)
+        assert ("model","decoder","layers","0","encoder_attn","q_proj","kernel") in flat
+        assert ("model","decoder","layers","0","encoder_attn_layer_norm","scale") in flat
+
+    def test_encoder_has_no_cross_attention(self, nllb_dir):
+        """Encoder tidak punya encoder_attn."""
+        params = waxjax.load(nllb_dir)
+        flat   = flatten_dict(params)
+        assert ("model","encoder","layers","0","encoder_attn","q_proj","kernel") not in flat
+
+    def test_lm_head_tied_to_shared(self, nllb_dir):
+        """lm_head.kernel == model.shared.embedding.T"""
+        params = waxjax.load(nllb_dir)
+        flat   = flatten_dict(params)
+        shared = np.array(flat[("model","shared","embedding")])  # (V,D)
+        lm     = np.array(flat[("lm_head","kernel")])            # (D,V)
+        assert lm.shape == (64, 512)
+        assert np.allclose(shared.T, lm, atol=1e-6)
+
+    def test_both_encoder_decoder_layers(self, nllb_dir):
+        params = waxjax.load(nllb_dir)
+        flat   = flatten_dict(params)
+        for i in range(2):
+            # Encoder
+            assert ("model","encoder","layers",str(i),"self_attn","q_proj","kernel") in flat
+            assert ("model","encoder","layers",str(i),"fc1","kernel") in flat
+            # Decoder
+            assert ("model","decoder","layers",str(i),"self_attn","q_proj","kernel") in flat
+            assert ("model","decoder","layers",str(i),"encoder_attn","q_proj","kernel") in flat

waxjax-0.1.0/tests/test_qwen2.py

@@ -0,0 +1,83 @@
+# tests/test_qwen2.py
+"""Integration test end-to-end untuk arsitektur Qwen2."""
+
+import pytest
+import numpy as np
+import jax.numpy as jnp
+from flax.traverse_util import flatten_dict
+import waxjax
+
+
+class TestQwen2Load:
+
+    def test_load_returns_nested_dict(self, qwen2_dir):
+        params = waxjax.load(qwen2_dir)
+        assert isinstance(params, dict)
+        assert "model" in params
+
+    def test_load_frozen_dict(self, qwen2_dir):
+        from flax.core import FrozenDict
+        params = waxjax.load(qwen2_dir, format="frozen")
+        assert isinstance(params, FrozenDict)
+
+    def test_embedding_shape(self, qwen2_dir):
+        params = waxjax.load(qwen2_dir)
+        embed  = np.array(params["model"]["embed_tokens"]["embedding"])
+        # (V=512, H=64) — tidak di-transpose
+        assert embed.shape == (512, 64)
+
+    def test_linear_kernel_transposed(self, qwen2_dir):
+        params = waxjax.load(qwen2_dir)
+        flat   = flatten_dict(params)
+        q_kernel = np.array(flat[("model","layers","0","self_attn","q_proj","kernel")])
+        # PT shape (QH*HD=64, H=64) → Flax shape (H=64, QH*HD=64)
+        assert q_kernel.shape == (64, 64)
+
+    def test_rmsnorm_has_scale_not_kernel(self, qwen2_dir):
+        params = waxjax.load(qwen2_dir)
+        flat   = flatten_dict(params)
+        assert ("model","layers","0","input_layernorm","scale") in flat
+        assert ("model","layers","0","input_layernorm","kernel") not in flat
+
+    def test_rmsnorm_no_bias(self, qwen2_dir):
+        """Qwen2 pakai RMSNorm, tidak ada bias."""
+        params = waxjax.load(qwen2_dir)
+        flat   = flatten_dict(params)
+        assert ("model","layers","0","input_layernorm","bias") not in flat
+
+    def test_tied_lm_head_created(self, qwen2_dir):
+        """lm_head.kernel harus ada meski tidak disimpan di file."""
+        params = waxjax.load(qwen2_dir)
+        flat   = flatten_dict(params)
+        assert ("lm_head","kernel") in flat
+        # Shape: (H=64, V=512)
+        assert flat[("lm_head","kernel")].shape == (64, 512)
+
+    def test_tied_lm_head_values_correct(self, qwen2_dir):
+        """lm_head.kernel == embed.T secara numerik."""
+        params = waxjax.load(qwen2_dir)
+        flat   = flatten_dict(params)
+        embed  = np.array(flat[("model","embed_tokens","embedding")])
+        lm     = np.array(flat[("lm_head","kernel")])
+        assert np.allclose(embed.T, lm, atol=1e-6)
+
+    def test_all_layers_present(self, qwen2_dir):
+        params = waxjax.load(qwen2_dir)
+        flat   = flatten_dict(params)
+        for i in range(2):  # NL=2 di fixture
+            assert ("model","layers",str(i),"self_attn","q_proj","kernel") in flat
+            assert ("model","layers",str(i),"mlp","gate_proj","kernel") in flat
+
+    def test_all_values_jnp_array(self, qwen2_dir):
+        params = waxjax.load(qwen2_dir)
+        flat   = flatten_dict(params)
+        for key, val in flat.items():
+            assert isinstance(val, jnp.ndarray), f"{key} bukan jnp.ndarray"
+
+    def test_no_nan_or_inf(self, qwen2_dir):
+        params = waxjax.load(qwen2_dir)
+        flat   = flatten_dict(params)
+        for key, val in flat.items():
+            arr = np.array(val)
+            assert not np.any(np.isnan(arr)), f"NaN di {key}"
+            assert not np.any(np.isinf(arr)), f"Inf di {key}"

waxjax-0.1.0/tests/test_rules.py

@@ -0,0 +1,77 @@
+# tests/test_rules.py
+"""Unit test untuk MappingRule — matching dan transform."""
+
+import pytest
+import numpy as np
+from waxjax.core.rules import MappingRule, transpose_2d, no_op
+
+
+class TestMappingRuleMatch:
+
+    def test_exact_match(self):
+        rule = MappingRule("lm_head.weight", rename=lambda k: k, transform="no_op")
+        assert rule.matches("lm_head.weight")
+        assert not rule.matches("lm_head.bias")
+
+    def test_single_wildcard(self):
+        rule = MappingRule(
+            "model.layers.*.self_attn.q_proj.weight",
+            rename=lambda k: k, transform="no_op"
+        )
+        assert rule.matches("model.layers.0.self_attn.q_proj.weight")
+        assert rule.matches("model.layers.11.self_attn.q_proj.weight")
+        # * tidak boleh match dot
+        assert not rule.matches("model.layers.0.1.self_attn.q_proj.weight")
+
+    def test_double_wildcard(self):
+        rule = MappingRule(
+            "model.**.weight",
+            rename=lambda k: k, transform="no_op"
+        )
+        assert rule.matches("model.layers.0.self_attn.q_proj.weight")
+        assert rule.matches("model.norm.weight")
+
+    def test_priority_auto(self):
+        specific = MappingRule(
+            "model.layers.*.self_attn_layer_norm.weight",
+            rename=lambda k: k, transform="no_op"
+        )
+        generic = MappingRule(
+            "model.layers.*.*.weight",
+            rename=lambda k: k, transform="no_op"
+        )
+        assert specific.priority > generic.priority
+
+    def test_invalid_transform_raises(self):
+        with pytest.raises(ValueError, match="tidak dikenal"):
+            MappingRule("*.weight", rename=lambda k: k, transform="unknown_transform")
+
+
+class TestMappingRuleTransform:
+
+    def test_transpose_2d(self):
+        arr = np.ones((4, 8))
+        rule = MappingRule("*.weight", rename=lambda k: k, transform="transpose_2d")
+        result = rule.apply_transform(arr)
+        assert result.shape == (8, 4)
+
+    def test_no_transpose_1d(self):
+        """1D array (bias) tidak boleh di-transpose meski pakai transpose_2d."""
+        arr = np.ones(8)
+        result = transpose_2d(arr)
+        assert result.shape == (8,)
+
+    def test_custom_callable_transform(self):
+        def double(arr):
+            return arr * 2
+        rule = MappingRule("*.weight", rename=lambda k: k, transform=double)
+        arr = np.array([1.0, 2.0])
+        assert np.allclose(rule.apply_transform(arr), [2.0, 4.0])
+
+    def test_rename_applied(self):
+        rule = MappingRule(
+            "*.weight",
+            rename=lambda k: k.replace(".weight", ".kernel"),
+            transform="no_op"
+        )
+        assert rule.apply_rename("model.fc.weight") == "model.fc.kernel"

waxjax-0.1.0/waxjax/__init__.py

@@ -0,0 +1,16 @@
+"""Compatibility shim: expose the public API under the expected
+package name `waxjax` while the real implementation lives in
+`safejax` (to preserve the original repository layout).
+"""
+# Re-export the public API without importing `safejax` package object
+# to avoid circular-import issues during package initialization.
+from .core.api import load, save, inspect, verify
+from .core.rules import MappingRule
+from .architectures._registry import register_architecture, ArchitectureConfig
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "load", "save", "inspect", "verify",
+    "MappingRule", "register_architecture", "ArchitectureConfig",
+]

waxjax-0.1.0/waxjax/architectures/__init__.py

@@ -0,0 +1,2 @@
+from . import _registry, qwen2, nllb
+__all__ = ["_registry", "qwen2", "nllb"]

waxjax-0.1.0/waxjax/architectures/_registry.py

@@ -0,0 +1 @@
+from safejax.architectures._registry import *

waxjax-0.1.0/waxjax/architectures/nllb.py

@@ -0,0 +1 @@
+from safejax.architectures.nllb import *

waxjax-0.1.0/waxjax/architectures/qwen2.py

@@ -0,0 +1 @@
+from safejax.architectures.qwen2 import *

waxjax-0.1.0/waxjax/core/__init__.py

@@ -0,0 +1,3 @@
+# waxjax.core package — wrappers re-exporting safejax.core
+from . import loader, mapper, nester, rules, api
+__all__ = ["loader", "mapper", "nester", "rules", "api"]

waxjax-0.1.0/waxjax/core/api.py

@@ -0,0 +1 @@
+from safejax.core.api import *

waxjax-0.1.0/waxjax/core/loader.py

@@ -0,0 +1 @@
+from safejax.core.loader import *

waxjax-0.1.0/waxjax/core/mapper.py

@@ -0,0 +1 @@
+from safejax.core.mapper import *

waxjax-0.1.0/waxjax/core/nester.py

@@ -0,0 +1 @@
+from safejax.core.nester import *

waxjax-0.1.0/waxjax/core/rules.py

@@ -0,0 +1 @@
+from safejax.core.rules import *

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

@@ -0,0 +1,55 @@
+Metadata-Version: 2.4
+Name: waxjax
+Version: 0.1.0
+Summary: Load HuggingFace SafeTensors models to JAX/Flax
+Home-page: https://github.com/ik12-b/waxjax
+Author: Your Name
+License: Apache-2.0
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: safetensors>=0.4.0
+Requires-Dist: numpy>=1.24
+Requires-Dist: jax>=0.4.0
+Provides-Extra: torch
+Requires-Dist: torch>=2.0; extra == "torch"
+Provides-Extra: flax
+Requires-Dist: flax>=0.7; extra == "flax"
+Provides-Extra: all
+Requires-Dist: torch>=2.0; extra == "all"
+Requires-Dist: flax>=0.7; extra == "all"
+Requires-Dist: ml_dtypes; extra == "all"
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: black; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Dynamic: author
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# waxjax
+
+Lightweight helper to load HuggingFace SafeTensors checkpoints into JAX/Flax.
+
+Install (latest from source):
+
+```bash
+pip install git+https://github.com/ik12-b/waxjax.git
+```
+
+Build and publish (recommended to test to TestPyPI first): see `RELEASE.md` or run locally:
+
+```bash
+python -m pip install --upgrade build twine
+python -m build
+twine upload --repository testpypi dist/*
+```

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

@@ -0,0 +1,27 @@
+MANIFEST.in
+README.md
+pyproject.toml
+setup.py
+safejax/__init__.py
+tests/test_api.py
+tests/test_loader.py
+tests/test_mapper.py
+tests/test_nllb.py
+tests/test_qwen2.py
+tests/test_rules.py
+waxjax/__init__.py
+waxjax.egg-info/PKG-INFO
+waxjax.egg-info/SOURCES.txt
+waxjax.egg-info/dependency_links.txt
+waxjax.egg-info/requires.txt
+waxjax.egg-info/top_level.txt
+waxjax/architectures/__init__.py
+waxjax/architectures/_registry.py
+waxjax/architectures/nllb.py
+waxjax/architectures/qwen2.py
+waxjax/core/__init__.py
+waxjax/core/api.py
+waxjax/core/loader.py
+waxjax/core/mapper.py
+waxjax/core/nester.py
+waxjax/core/rules.py

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

@@ -0,0 +1 @@
+

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

@@ -0,0 +1,19 @@
+safetensors>=0.4.0
+numpy>=1.24
+jax>=0.4.0
+
+[all]
+torch>=2.0
+flax>=0.7
+ml_dtypes
+
+[dev]
+pytest
+black
+ruff
+
+[flax]
+flax>=0.7
+
+[torch]
+torch>=2.0

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

@@ -0,0 +1,2 @@
+safejax
+waxjax