pytorchmlx 0.0.1__tar.gz

pytorchmlx-0.0.1/.github/workflows/publish.yml

@@ -0,0 +1,110 @@
+name: publish to pypi
+
+on:
+  workflow_dispatch:
+
+concurrency:
+  group: publish-to-pypi
+  cancel-in-progress: false
+
+permissions:
+  contents: write
+
+jobs:
+  publish:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: install uv
+        uses: astral-sh/setup-uv@v6
+
+      - name: select release
+        id: release
+        env:
+          GH_TOKEN: ${{ github.token }}
+        run: |
+          latest_tag=$(git tag --list 'v[0-9]*' --sort=-v:refname | head -n 1)
+          create_tag=false
+          publish=true
+
+          if [ -z "$latest_tag" ]; then
+            target_tag=v0.0.1
+            create_tag=true
+          elif gh release view "$latest_tag" >/dev/null 2>&1; then
+            if git diff --quiet "$latest_tag"..HEAD; then
+              publish=false
+              target_tag="$latest_tag"
+            else
+              version=${latest_tag#v}
+              IFS=. read -r major minor patch <<< "$version"
+              target_tag="v${major}.${minor}.$((patch + 1))"
+              create_tag=true
+            fi
+          else
+            target_tag="$latest_tag"
+          fi
+
+          if [[ ! "$target_tag" =~ ^v[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
+            echo "invalid release tag: $target_tag"
+            exit 1
+          fi
+
+          {
+            echo "publish=$publish"
+            echo "create_tag=$create_tag"
+            echo "tag=$target_tag"
+            echo "version=${target_tag#v}"
+          } >> "$GITHUB_OUTPUT"
+
+      - name: stop when unchanged
+        if: steps.release.outputs.publish == 'false'
+        run: echo "no changes since ${{ steps.release.outputs.tag }}"
+
+      - name: verify pypi credentials
+        if: steps.release.outputs.publish == 'true'
+        env:
+          PYPI_API_TOKEN: ${{ secrets.PYPI_API_TOKEN }}
+        run: |
+          if [ -z "$PYPI_API_TOKEN" ]; then
+            echo "PYPI_API_TOKEN is not configured"
+            exit 1
+          fi
+
+      - name: set package version
+        if: steps.release.outputs.publish == 'true'
+        env:
+          VERSION: ${{ steps.release.outputs.version }}
+        run: sed -i "0,/version = \".*\"/s//version = \"$VERSION\"/" pyproject.toml
+
+      - name: build package
+        if: steps.release.outputs.publish == 'true'
+        run: uv build
+
+      - name: check package
+        if: steps.release.outputs.publish == 'true'
+        run: uv tool run twine check dist/*
+
+      - name: create tag
+        if: steps.release.outputs.publish == 'true' && steps.release.outputs.create_tag == 'true'
+        env:
+          TAG: ${{ steps.release.outputs.tag }}
+        run: |
+          git tag "$TAG" "$GITHUB_SHA"
+          git push origin "$TAG"
+
+      - name: publish package
+        if: steps.release.outputs.publish == 'true'
+        env:
+          TWINE_USERNAME: __token__
+          TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
+        run: uv tool run twine upload --verbose dist/*
+
+      - name: create github release
+        if: steps.release.outputs.publish == 'true'
+        env:
+          GH_TOKEN: ${{ github.token }}
+          TAG: ${{ steps.release.outputs.tag }}
+        run: gh release create "$TAG" dist/* --generate-notes --title "$TAG"

pytorchmlx-0.0.1/.gitignore

@@ -0,0 +1,14 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+examples/tinystories-llm/TinyStoriesV2-GPT4-valid.txt
+
+# Virtual environments
+.venv
+
+# Environment files
+.env

pytorchmlx-0.0.1/.python-version

@@ -0,0 +1 @@
+3.13

pytorchmlx-0.0.1/LICENSE

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

pytorchmlx-0.0.1/PKG-INFO

@@ -0,0 +1,48 @@
+Metadata-Version: 2.5
+Name: pytorchmlx
+Version: 0.0.1
+Summary: An educational PyTorch-shaped interface for MLX and PyTorch
+Author: Priyanshu Jain
+License: MIT License
+        
+        Copyright (c) 2026 Priyanshu Jain
+        
+        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.10
+Requires-Dist: mlx<0.33,>=0.32.2
+Requires-Dist: numpy
+Requires-Dist: torch<3,>=2.4
+Description-Content-Type: text/markdown
+
+# torchmlx
+
+torchmlx is a pytorch-shaped compatibility layer that uses mlx on apple silicon and pytorch elsewhere.
+
+NOTE: it is experimental and built first for educational use. If you find a bug, please create an issue on the repo.
+
+```python
+import torchmlx as torch
+from torchmlx import nn, optim
+
+model = nn.Linear(4, 2)
+optimizer = optim.AdamW(model.parameters(), lr=3e-4)
+```
+
+see the [tinystories example](examples/tinystories-llm/train.py) and [compatibility details](docs/compatibility.md).

pytorchmlx-0.0.1/README.md

@@ -0,0 +1,15 @@
+# torchmlx
+
+torchmlx is a pytorch-shaped compatibility layer that uses mlx on apple silicon and pytorch elsewhere.
+
+NOTE: it is experimental and built first for educational use. If you find a bug, please create an issue on the repo.
+
+```python
+import torchmlx as torch
+from torchmlx import nn, optim
+
+model = nn.Linear(4, 2)
+optimizer = optim.AdamW(model.parameters(), lr=3e-4)
+```
+
+see the [tinystories example](examples/tinystories-llm/train.py) and [compatibility details](docs/compatibility.md).

pytorchmlx-0.0.1/docs/compatibility.md

@@ -0,0 +1,25 @@
+# Compatibility
+
+TorchMLX targets the common transformer operations used by GPT-2, Llama 3, Qwen 3, and GPT-OSS style implementations.
+
+Supported MLX operations include embeddings, linear layers, normalization building blocks, dropout, activations, causal attention, tensor shape operations, masks, top-k routing, and AdamW training.
+
+MLX arrays remain native arrays. Torch-style tensor methods are installed on the native array type for the supported subset.
+
+Boolean expert routing and `unique` execute eagerly because their output shapes control Python flow.
+
+The standard training sequence works with `cross_entropy` losses applied directly or after reshape, view, slicing, transpose, squeeze, unsqueeze, or flatten operations:
+
+```python
+optimizer.zero_grad()
+logits = model(input_tokens)
+loss = F.cross_entropy(logits.reshape(-1, vocabulary_size), targets.reshape(-1))
+loss.backward()
+optimizer.step()
+```
+
+MLX implements this sequence by recording the outer model call and replaying it inside `value_and_grad` during `backward`. Random state is restored for the replay so dropout uses the same mask. Unrecorded loss expressions, gradient hooks, parameter `.grad`, higher-order gradients, and multiple-forward losses remain unsupported.
+
+Set `TORCHMLX_BACKEND=torch` before import to use native PyTorch for unsupported programs. TorchMLX never changes backend during an operation.
+
+The referenced OpenArch model files contain source errors independent of TorchMLX, including invalid constructor calls and undefined attributes. Correct those errors before using either backend.

pytorchmlx-0.0.1/examples/tinystories-llm/train.py

@@ -0,0 +1,263 @@
+import random
+import urllib.request
+from dataclasses import dataclass
+from pathlib import Path
+
+import torchmlx as torch
+from torchmlx import nn, optim
+import torchmlx.nn.functional as F
+
+
+device = torch.device(
+    "mps"
+    if torch.backends.mps.is_available()
+    else "cuda"
+    if torch.cuda.is_available()
+    else "cpu"
+)
+
+DATA_URL = "https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStoriesV2-GPT4-valid.txt"
+DATA_PATH = Path(__file__).with_name("TinyStoriesV2-GPT4-valid.txt")
+
+
+@dataclass
+class TransformerConfig:
+    vocabulary_size: int
+    context_length: int = 128
+    model_dimension: int = 32
+    head_count: int = 16
+    layer_count: int = 4
+    feed_forward_dimension: int = 128
+
+
+@dataclass
+class TrainingConfig:
+    batch_size: int = 32
+    step_count: int = 5000
+    learning_rate: float = 1e-3
+    weight_decay: float = 1e-2
+    log_interval: int = 100
+
+
+class CharacterTokenizer:
+    def __init__(self, text):
+        self.characters = sorted(set(text))
+        self.token_by_character = {
+            character: token for token, character in enumerate(self.characters)
+        }
+
+    def encode(self, text):
+        return [self.token_by_character[character] for character in text]
+
+    def decode(self, tokens):
+        return "".join(self.characters[int(token)] for token in tokens)
+
+    def __len__(self):
+        return len(self.characters)
+
+
+class TokenEmbedding(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.embedding = nn.Embedding(
+            config.vocabulary_size, config.model_dimension
+        )
+
+    def forward(self, tokens):
+        return self.embedding(tokens)
+
+
+class PositionalEmbedding(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.embedding = nn.Embedding(
+            config.context_length, config.model_dimension
+        )
+
+    def forward(self, tokens):
+        positions = torch.arange(tokens.shape[1], device=tokens.device)
+        return self.embedding(positions)
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.head_count = config.head_count
+        self.head_dimension = config.model_dimension // config.head_count
+        self.query = nn.Linear(config.model_dimension, config.model_dimension)
+        self.key = nn.Linear(config.model_dimension, config.model_dimension)
+        self.value = nn.Linear(config.model_dimension, config.model_dimension)
+        self.output = nn.Linear(config.model_dimension, config.model_dimension)
+
+    def split_heads(self, hidden_states):
+        batch_size, sequence_length, _ = hidden_states.shape
+        hidden_states = hidden_states.reshape(
+            batch_size, sequence_length, self.head_count, self.head_dimension
+        )
+        return hidden_states.transpose(1, 2)
+
+    def merge_heads(self, hidden_states):
+        batch_size, _, sequence_length, _ = hidden_states.shape
+        hidden_states = hidden_states.transpose(1, 2).contiguous()
+        return hidden_states.reshape(
+            batch_size, sequence_length, self.head_count * self.head_dimension
+        )
+
+    def forward(self, hidden_states):
+        queries = self.split_heads(self.query(hidden_states))
+        keys = self.split_heads(self.key(hidden_states))
+        values = self.split_heads(self.value(hidden_states))
+
+        attention_scores = torch.matmul(queries, keys.transpose(-2, -1))
+        attention_scores = attention_scores / self.head_dimension**0.5
+
+        sequence_length = hidden_states.shape[1]
+        future_positions = torch.triu(
+            torch.ones(
+                sequence_length,
+                sequence_length,
+                dtype=torch.bool,
+                device=hidden_states.device,
+            ),
+            diagonal=1,
+        )
+        attention_scores = attention_scores.masked_fill(
+            future_positions, float("-inf")
+        )
+        attention_weights = F.softmax(attention_scores, dim=-1)
+        attended_values = torch.matmul(attention_weights, values)
+        return self.output(self.merge_heads(attended_values))
+
+
+class FeedForward(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.layers = nn.Sequential(
+            nn.Linear(config.model_dimension, config.feed_forward_dimension),
+            nn.GELU(),
+            nn.Linear(config.feed_forward_dimension, config.model_dimension),
+        )
+
+    def forward(self, hidden_states):
+        return self.layers(hidden_states)
+
+
+class Unembedding(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.output = nn.Linear(
+            config.model_dimension, config.vocabulary_size, bias=False
+        )
+
+    def forward(self, hidden_states):
+        return self.output(hidden_states)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.attention_norm = nn.LayerNorm(config.model_dimension)
+        self.attention = CausalSelfAttention(config)
+        self.feed_forward_norm = nn.LayerNorm(config.model_dimension)
+        self.feed_forward = FeedForward(config)
+
+    def forward(self, hidden_states):
+        hidden_states = hidden_states + self.attention(
+            self.attention_norm(hidden_states)
+        )
+        return hidden_states + self.feed_forward(
+            self.feed_forward_norm(hidden_states)
+        )
+
+
+class TinyStoriesTransformer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.token_embedding = TokenEmbedding(config)
+        self.position_embedding = PositionalEmbedding(config)
+        self.blocks = nn.ModuleList(
+            [TransformerBlock(config) for _ in range(config.layer_count)]
+        )
+        self.final_norm = nn.LayerNorm(config.model_dimension)
+        self.unembedding = Unembedding(config)
+
+    def forward(self, tokens):
+        hidden_states = self.token_embedding(tokens)
+        hidden_states = hidden_states + self.position_embedding(tokens)
+        for block in self.blocks:
+            hidden_states = block(hidden_states)
+        return self.unembedding(self.final_norm(hidden_states))
+
+
+def load_tiny_stories():
+    if not DATA_PATH.exists():
+        print(f"downloading TinyStories to {DATA_PATH}")
+        urllib.request.urlretrieve(DATA_URL, DATA_PATH)
+    return DATA_PATH.read_text(encoding="utf-8")
+
+
+def create_batch(encoded_text, batch_size, context_length):
+    starts = [
+        random.randrange(len(encoded_text) - context_length - 1)
+        for _ in range(batch_size)
+    ]
+    input_tokens = [
+        encoded_text[start : start + context_length] for start in starts
+    ]
+    target_tokens = [
+        encoded_text[start + 1 : start + context_length + 1]
+        for start in starts
+    ]
+    return torch.tensor(input_tokens, dtype=torch.long, device=device), torch.tensor(
+        target_tokens, dtype=torch.long, device=device
+    )
+
+
+def language_model_loss(logits, target_tokens):
+    vocabulary_size = logits.shape[-1]
+    return F.cross_entropy(
+        logits.reshape(-1, vocabulary_size), target_tokens.reshape(-1)
+    )
+
+
+def generate_text(model, tokenizer, prompt, token_count):
+    generated_tokens = tokenizer.encode(prompt)
+    model.eval()
+    for _ in range(token_count):
+        context = generated_tokens[-model.config.context_length :]
+        input_tokens = torch.tensor([context], dtype=torch.long, device=device)
+        next_token_logits = model(input_tokens)[:, -1, :]
+        next_token = torch.categorical(next_token_logits).item()
+        generated_tokens.append(next_token)
+    return tokenizer.decode(generated_tokens)
+
+
+training_text = load_tiny_stories()
+tokenizer = CharacterTokenizer(training_text)
+config = TransformerConfig(vocabulary_size=len(tokenizer))
+training_config = TrainingConfig()
+encoded_text = tokenizer.encode(training_text)
+model = TinyStoriesTransformer(config).to(device)
+optimizer = optim.AdamW(
+    model.parameters(),
+    lr=training_config.learning_rate,
+    weight_decay=training_config.weight_decay,
+)
+model.train()
+
+for step in range(training_config.step_count):
+    input_tokens, target_tokens = create_batch(
+        encoded_text,
+        batch_size=training_config.batch_size,
+        context_length=config.context_length,
+    )
+    optimizer.zero_grad()
+    logits = model(input_tokens)
+    loss = language_model_loss(logits, target_tokens)
+    loss.backward()
+    optimizer.step()
+    if step % training_config.log_interval == 0:
+        print(f"step {step}: loss {loss.item():.4f}")
+
+print(generate_text(model, tokenizer, "Once upon a time", token_count=120))

pytorchmlx-0.0.1/pyproject.toml

@@ -0,0 +1,22 @@
+[project]
+name = "pytorchmlx"
+version = "0.0.1"
+description = "An educational PyTorch-shaped interface for MLX and PyTorch"
+readme = "README.md"
+requires-python = ">=3.10"
+license = { file = "LICENSE" }
+authors = [
+    { name = "Priyanshu Jain" },
+]
+dependencies = [
+    "mlx>=0.32.2,<0.33",
+    "numpy",
+    "torch>=2.4,<3",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/torchmlx"]