pytorchmlx 0.0.1__py3-none-any.whl

pytorchmlx-0.0.1.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,12 @@
+torchmlx/__init__.py,sha256=Vo3_qG5ib3rDOKUT1N6y7Fr4760MM_yecWUirIBPunQ,8851
+torchmlx/_autograd.py,sha256=HkkesMlbAvFnHc47AenA4Sapp2A7jpXSIOX9u5XJoXc,3866
+torchmlx/_backend.py,sha256=EPBqEX402BLx1DjZ7Ay0N7IAds85GxVZQgSOdZn0qyE,720
+torchmlx/_mlx_tensor.py,sha256=ePvQltx5Oh9_vDFetO5AnMxE58h9ViwBzo_z6GRyZv4,6936
+torchmlx/trainer.py,sha256=CQeNKQfZu-DZ6bQ8InTAQZBAZMOS96HEitW-q8J3J8M,1837
+torchmlx/nn/__init__.py,sha256=QKw7NH8dWyovShPY-SOnZrk3nGG5Gh_DXreSUuIXQr4,8510
+torchmlx/nn/functional.py,sha256=PuGF28lapTv_68x7ejvNoOPjGptcWNWijMoiCi3Zpnw,4457
+torchmlx/optim/__init__.py,sha256=nSX-wNGAfz6DaGzuNVSEnGFuzq25OrfmdGguM5OQnek,2088
+pytorchmlx-0.0.1.dist-info/METADATA,sha256=w3SNzTxmyQ0bmcTzyjtXveFux17N6k-Z_DiTnDmE5bw,2056
+pytorchmlx-0.0.1.dist-info/WHEEL,sha256=THafob7ofN-NsuMN7Mg4qZyHaQI7KkD-QlcQatYhXPo,87
+pytorchmlx-0.0.1.dist-info/licenses/LICENSE,sha256=Rc24jdpcK087K5-G3PHFZ8AFoZ32KT95EyqvTMvWAQQ,1071
+pytorchmlx-0.0.1.dist-info/RECORD,,

pytorchmlx-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.32.3
+Root-Is-Purelib: true
+Tag: py3-none-any

pytorchmlx-0.0.1.dist-info/licenses/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.

torchmlx/__init__.py

@@ -0,0 +1,305 @@
+import builtins as _builtins
+import math as _math
+import platform as _platform
+from types import SimpleNamespace as _SimpleNamespace
+
+from ._backend import BACKEND, unsupported
+
+
+def current_backend():
+    return BACKEND
+
+
+if BACKEND == "torch":
+    import torch as _native
+
+    Tensor = _native.Tensor
+    tensor = _native.tensor
+    from_numpy = _native.from_numpy
+    arange = _native.arange
+    randint = _native.randint
+    chunk = _native.chunk
+    transpose = _native.transpose
+    float16 = _native.float16
+    float32 = _native.float32
+    float64 = _native.float64
+    bfloat16 = _native.bfloat16
+    int8 = _native.int8
+    int16 = _native.int16
+    int32 = _native.int32
+    int64 = _native.int64
+    uint8 = _native.uint8
+    bool = _native.bool
+    half = float16
+    float = float32
+    double = float64
+    int = int32
+    long = int64
+
+    def categorical(logits, dim=-1, num_samples=1):
+        probabilities = _native.softmax(logits, dim=dim)
+        return _native.multinomial(probabilities, num_samples=num_samples)
+
+    def __getattr__(name):
+        return getattr(_native, name)
+
+else:
+    import mlx.core as _native
+    import numpy as _np
+
+    Tensor = _native.array
+    float16 = _native.float16
+    float32 = _native.float32
+    float64 = _native.float64
+    bfloat16 = _native.bfloat16
+    int8 = _native.int8
+    int16 = _native.int16
+    int32 = _native.int32
+    int64 = _native.int64
+    uint8 = _native.uint8
+    bool = _native.bool_
+    half = float16
+    float = float32
+    double = float64
+    int = int32
+    long = int64
+
+    pi = _math.pi
+
+    class device:
+        def __init__(self, value):
+            if isinstance(value, device):
+                value = value.type
+            value = str(value)
+            if value != "mps":
+                raise ValueError("the MLX backend only accepts device='mps'")
+            self.type = value
+            self.index = None
+
+        def __str__(self):
+            return self.type
+
+        def __repr__(self):
+            return f"device(type={self.type!r})"
+
+        def __eq__(self, other):
+            return str(other) == self.type
+
+    class _MPS:
+        @staticmethod
+        def is_available():
+            return _platform.system() == "Darwin" and _platform.machine() == "arm64"
+
+    class _CUDA:
+        @staticmethod
+        def is_available():
+            return False
+
+    backends = _SimpleNamespace(mps=_MPS())
+    cuda = _CUDA()
+
+    from ._mlx_tensor import install as _install_tensor_methods
+
+    _install_tensor_methods(device)
+
+    def _check_device(device):
+        if device is not None and str(device) != "mps":
+            raise ValueError("the MLX backend only accepts device='mps'")
+
+    def tensor(data, dtype=None, device=None, requires_grad=False, pin_memory=False):
+        _check_device(device)
+        if requires_grad:
+            raise RuntimeError(
+                "requires_grad is not supported by the MLX backend; use torchmlx.Trainer"
+            )
+        if pin_memory:
+            raise RuntimeError("pin_memory is not supported by the MLX backend")
+        if dtype is None and not isinstance(data, (_native.array, _np.ndarray)):
+            kind = _np.asarray(data).dtype.kind
+            if kind in {"i", "u"}:
+                dtype = int64
+            elif kind == "b":
+                dtype = bool
+        return _native.array(data, dtype=dtype)
+
+    def from_numpy(array):
+        if not isinstance(array, _np.ndarray):
+            raise TypeError("from_numpy expects a numpy.ndarray")
+        return _native.array(array)
+
+    def arange(start, end=None, step=1, *, dtype=None, device=None):
+        _check_device(device)
+        if end is None:
+            start, end = 0, start
+        if dtype is None and all(
+            isinstance(value, _builtins.int) for value in (start, end, step)
+        ):
+            dtype = int64
+        return _native.arange(start, end, step, dtype=dtype)
+
+    def randint(low, high, size, *, dtype=None, device=None, requires_grad=False):
+        _check_device(device)
+        if requires_grad:
+            raise RuntimeError(
+                "requires_grad is not supported by the MLX backend; use torchmlx.Trainer"
+            )
+        result = _native.random.randint(low, high, shape=size)
+        return result.astype(dtype or int64)
+
+    def chunk(input, chunks, dim=0):
+        return input.chunk(chunks, dim=dim)
+
+    def transpose(input, dim0, dim1):
+        return input.transpose(dim0, dim1)
+
+    def zeros(*size, dtype=None, device=None, requires_grad=False):
+        _check_device(device)
+        if requires_grad:
+            raise RuntimeError(
+                "requires_grad is not supported by the MLX backend; use torchmlx.Trainer"
+            )
+        shape = size[0] if len(size) == 1 and isinstance(size[0], (tuple, list)) else size
+        return _native.zeros(shape, dtype=dtype or float32)
+
+    def ones(*size, dtype=None, device=None, requires_grad=False):
+        _check_device(device)
+        if requires_grad:
+            raise RuntimeError(
+                "requires_grad is not supported by the MLX backend; use torchmlx.Trainer"
+            )
+        shape = size[0] if len(size) == 1 and isinstance(size[0], (tuple, list)) else size
+        return _native.ones(shape, dtype=dtype or float32)
+
+    def zeros_like(input, *, dtype=None, device=None, requires_grad=False):
+        _check_device(device)
+        if requires_grad:
+            raise RuntimeError(
+                "requires_grad is not supported by the MLX backend; use torchmlx.Trainer"
+            )
+        return _native.zeros_like(input).astype(dtype or input.dtype)
+
+    def ones_like(input, *, dtype=None, device=None, requires_grad=False):
+        _check_device(device)
+        if requires_grad:
+            raise RuntimeError(
+                "requires_grad is not supported by the MLX backend; use torchmlx.Trainer"
+            )
+        return _native.ones_like(input).astype(dtype or input.dtype)
+
+    def cat(tensors, dim=0):
+        return _native.concatenate(tensors, axis=dim)
+
+    def softmax(input, dim, dtype=None):
+        value = input.astype(dtype) if dtype is not None else input
+        return _native.softmax(value, axis=dim)
+
+    def triu(input, diagonal=0):
+        return _native.triu(input, k=diagonal)
+
+    def tril(input, diagonal=0):
+        return _native.tril(input, k=diagonal)
+
+    def topk(input, k, dim=None, largest=True, sorted=True):
+        axis = -1 if dim is None else dim
+        indices = _native.argsort(input, axis=axis)
+        indices = _native.flip(indices, axis=axis) if largest else indices
+        slices = [slice(None)] * input.ndim
+        slices[axis] = slice(0, k)
+        indices = indices[tuple(slices)].astype(int64)
+        values = _native.take_along_axis(input, indices, axis=axis)
+        return values, indices
+
+    def unique(input, sorted=True, return_inverse=False, return_counts=False, dim=None):
+        if return_inverse or return_counts or dim is not None:
+            unsupported("torchmlx.unique with non-default options")
+        values = _native.sort(input.reshape(-1))
+        if values.shape[0] < 2:
+            return values
+        keep = _native.concatenate(
+            [_native.array([True]), values[1:] != values[:-1]]
+        )
+        return values[keep]
+
+    exp = _native.exp
+    sin = _native.sin
+    cos = _native.cos
+    tanh = _native.tanh
+    sqrt = _native.sqrt
+    matmul = _native.matmul
+    outer = _native.outer
+
+    def pow(input, exponent):
+        return _native.power(input, exponent)
+
+    def polar(abs, angle):
+        return abs * _native.exp(_native.array(1j) * angle)
+
+    def categorical(logits, dim=-1, num_samples=1):
+        if num_samples == 1:
+            return _native.random.categorical(logits, axis=dim)[..., None].astype(int64)
+        return _native.random.categorical(
+            logits, axis=dim, num_samples=num_samples
+        ).astype(int64)
+
+    def __getattr__(name):
+        unsupported(f"torchmlx.{name}")
+
+
+import importlib as _importlib
+
+nn = _importlib.import_module("torchmlx.nn")
+optim = _importlib.import_module("torchmlx.optim")
+from .trainer import Trainer
+
+__all__ = [
+    "Tensor",
+    "Trainer",
+    "arange",
+    "backends",
+    "bfloat16",
+    "bool",
+    "categorical",
+    "cat",
+    "chunk",
+    "cuda",
+    "current_backend",
+    "device",
+    "double",
+    "float",
+    "float16",
+    "float32",
+    "float64",
+    "exp",
+    "from_numpy",
+    "int8",
+    "int16",
+    "int32",
+    "int64",
+    "half",
+    "int",
+    "long",
+    "matmul",
+    "nn",
+    "ones",
+    "ones_like",
+    "optim",
+    "outer",
+    "pi",
+    "polar",
+    "pow",
+    "randint",
+    "tensor",
+    "softmax",
+    "sin",
+    "cos",
+    "sqrt",
+    "tanh",
+    "topk",
+    "transpose",
+    "tril",
+    "triu",
+    "uint8",
+    "unique",
+    "zeros",
+    "zeros_like",
+]

torchmlx/_autograd.py

@@ -0,0 +1,135 @@
+import mlx.core as mx
+import mlx.nn as nn
+
+
+_active_optimizer = None
+_forward_depth = 0
+_latest_forward = None
+_loss_plans = {}
+_lineage = {}
+_replayed_losses = {}
+_suspended = False
+
+
+def _snapshot_random_state():
+    state = [key + mx.array(0, dtype=key.dtype) for key in mx.random.state]
+    mx.eval(state)
+    return state
+
+
+def _restore_random_state(state):
+    for current_key, saved_key in zip(mx.random.state, state):
+        current_key[...] = saved_key
+    mx.eval(mx.random.state)
+
+
+def begin_forward():
+    global _forward_depth
+    _forward_depth += 1
+
+
+def end_forward(model, args, kwargs, output):
+    global _forward_depth, _latest_forward
+    _forward_depth -= 1
+    if (
+        _forward_depth == 0
+        and _active_optimizer is not None
+        and model is _active_optimizer._model
+        and not _suspended
+        and isinstance(output, mx.array)
+    ):
+        _latest_forward = (model, args, kwargs, lambda value: value)
+        _lineage[id(output)] = (output, _latest_forward)
+
+
+def abort_forward():
+    global _forward_depth
+    _forward_depth -= 1
+
+
+def activate(optimizer):
+    global _active_optimizer, _latest_forward
+    _active_optimizer = optimizer
+    _latest_forward = None
+    _loss_plans.clear()
+    _lineage.clear()
+    _replayed_losses.clear()
+    optimizer._random_before_forward = _snapshot_random_state()
+
+
+def propagate(source, result, operation=lambda value: value):
+    entry = _lineage.get(id(source))
+    if not _suspended and entry is not None and entry[0] is source:
+        model, args, kwargs, previous = entry[1]
+        _lineage[id(result)] = (
+            result,
+            (
+                model,
+                args,
+                kwargs,
+                lambda output: operation(previous(output)),
+            ),
+        )
+    return result
+
+
+def register_loss(loss, loss_input, rebuild):
+    if _suspended or _active_optimizer is None:
+        return loss
+    entry = _lineage.get(id(loss_input))
+    if entry is None or entry[0] is not loss_input:
+        return loss
+    model, args, kwargs, transform = entry[1]
+
+    def plan():
+        def objective():
+            return rebuild(transform(model(*args, **kwargs)))
+
+        return nn.value_and_grad(model, objective)()
+
+    _loss_plans[id(loss)] = (loss, model, plan)
+    return loss
+
+
+def backward(loss):
+    global _suspended
+    if _active_optimizer is None:
+        raise RuntimeError("optimizer.zero_grad() must be called before loss.backward()")
+    entry = _loss_plans.get(id(loss))
+    if entry is None or entry[0] is not loss:
+        raise RuntimeError(
+            "this MLX loss cannot use backward compatibility; compute it with a supported torchmlx loss function"
+        )
+    _, model, plan = entry
+    random_after_forward = _snapshot_random_state()
+    _restore_random_state(_active_optimizer._random_before_forward)
+    _suspended = True
+    try:
+        replayed_loss, gradients = plan()
+        mx.eval(replayed_loss, gradients, mx.random.state)
+    finally:
+        _suspended = False
+        _restore_random_state(random_after_forward)
+    _replayed_losses[id(loss)] = (loss, replayed_loss)
+    _active_optimizer._pending_update = (model, gradients)
+
+
+def replayed_value(value):
+    entry = _replayed_losses.get(id(value))
+    if entry is None or entry[0] is not value:
+        return value
+    return entry[1]
+
+
+def step(optimizer):
+    global _active_optimizer, _latest_forward
+    if optimizer._pending_update is None:
+        raise RuntimeError("loss.backward() must be called before optimizer.step()")
+    model, gradients = optimizer._pending_update
+    optimizer._optimizer.update(model, gradients)
+    mx.eval(model.parameters(), optimizer._optimizer.state, mx.random.state)
+    optimizer._pending_update = None
+    _active_optimizer = None
+    _latest_forward = None
+    _loss_plans.clear()
+    _lineage.clear()

torchmlx/_backend.py

@@ -0,0 +1,27 @@
+import os
+import platform
+
+
+def _select_backend():
+    override = os.environ.get("TORCHMLX_BACKEND")
+    if override is not None:
+        backend = override.lower()
+        if backend not in {"mlx", "torch"}:
+            raise RuntimeError(
+                "TORCHMLX_BACKEND must be either 'mlx' or 'torch', "
+                f"not {override!r}"
+            )
+        return backend
+    if platform.system() == "Darwin" and platform.machine() == "arm64":
+        return "mlx"
+    return "torch"
+
+
+BACKEND = _select_backend()
+
+
+def unsupported(api):
+    raise NotImplementedError(
+        f"{api} is not supported by the MLX backend. "
+        "Set TORCHMLX_BACKEND=torch before importing torchmlx to use PyTorch."
+    )

torchmlx/_mlx_tensor.py

@@ -0,0 +1,240 @@
+import mlx.core as mx
+
+
+_transpose = mx.array.transpose
+_reshape = mx.array.reshape
+_squeeze = mx.array.squeeze
+_mean = mx.array.mean
+_var = mx.array.var
+_any = mx.array.any
+_getitem = mx.array.__getitem__
+_setitem = mx.array.__setitem__
+_item = mx.array.item
+
+
+def _torch_transpose(self, dim0=None, dim1=None):
+    if dim1 is None:
+        if dim0 is None:
+            axes = list(reversed(range(self.ndim)))
+        else:
+            axes = dim0
+    else:
+        axes = list(range(self.ndim))
+        axes[dim0], axes[dim1] = axes[dim1], axes[dim0]
+    result = _transpose(self, axes)
+    from ._autograd import propagate
+
+    return propagate(self, result, lambda value: _transpose(value, axes))
+
+
+def _view(self, *shape):
+    if len(shape) == 1 and isinstance(shape[0], (tuple, list)):
+        shape = shape[0]
+    return _torch_reshape(self, shape)
+
+
+def _torch_reshape(self, *shape):
+    if len(shape) == 1 and isinstance(shape[0], (tuple, list)):
+        shape = shape[0]
+    result = _reshape(self, shape)
+    from ._autograd import propagate
+
+    return propagate(self, result, lambda value: _reshape(value, shape))
+
+
+def _unsqueeze(self, dim):
+    result = mx.expand_dims(self, axis=dim)
+    from ._autograd import propagate
+
+    return propagate(self, result, lambda value: mx.expand_dims(value, axis=dim))
+
+
+def _torch_squeeze(self, dim=None):
+    result = _squeeze(self, axis=dim)
+    from ._autograd import propagate
+
+    return propagate(self, result, lambda value: _squeeze(value, axis=dim))
+
+
+def _flatten(self, start_dim=0, end_dim=-1):
+    if end_dim < 0:
+        end_dim += self.ndim
+    flattened = 1
+    for dimension in self.shape[start_dim : end_dim + 1]:
+        flattened *= dimension
+    shape = self.shape[:start_dim] + (flattened,) + self.shape[end_dim + 1 :]
+    return _torch_reshape(self, shape)
+
+
+def _float(self):
+    return self.astype(mx.float32)
+
+
+def _bool(self):
+    return self.astype(mx.bool_)
+
+
+def _pow(self, exponent):
+    return mx.power(self, exponent)
+
+
+def _torch_mean(self, dim=None, keepdim=False, dtype=None):
+    value = self.astype(dtype) if dtype is not None else self
+    return _mean(value, axis=dim, keepdims=keepdim)
+
+
+def _torch_var(
+    self,
+    dim=None,
+    unbiased=True,
+    keepdim=False,
+    *,
+    correction=None,
+):
+    ddof = int(unbiased) if correction is None else correction
+    return _var(self, axis=dim, keepdims=keepdim, ddof=ddof)
+
+
+def _size(self, dim=None):
+    return self.shape if dim is None else self.shape[dim]
+
+
+def _chunk(self, chunks, dim=0):
+    if chunks <= 0:
+        raise ValueError("chunks must be greater than 0")
+    length = self.shape[dim]
+    if length == 0:
+        return tuple(mx.split(self, chunks, axis=dim))
+    chunk_size = (length + chunks - 1) // chunks
+    indices = list(range(chunk_size, length, chunk_size))
+    return tuple(mx.split(self, indices, axis=dim))
+
+
+def _to(self, *args, dtype=None, device=None, **kwargs):
+    if kwargs:
+        name = next(iter(kwargs))
+        raise TypeError(f"to() got an unexpected keyword argument {name!r}")
+    for value in args:
+        if isinstance(value, mx.Dtype):
+            dtype = value
+        elif isinstance(value, mx.array):
+            dtype = value.dtype
+        else:
+            device = value
+    if device is not None and str(device) != "mps":
+        raise ValueError("the MLX backend only accepts device='mps'")
+    return self.astype(dtype) if dtype is not None and dtype != self.dtype else self
+
+
+def _repeat_interleave(self, repeats, dim=None):
+    return mx.repeat(self, repeats, axis=dim)
+
+
+def _masked_fill(self, mask, value):
+    return mx.where(mask, mx.array(value, dtype=self.dtype), self)
+
+
+def _expand(self, *sizes):
+    if len(sizes) == 1 and isinstance(sizes[0], (tuple, list)):
+        sizes = tuple(sizes[0])
+    if len(sizes) < self.ndim:
+        raise ValueError("expanded size must have at least as many dimensions as the tensor")
+    source = (1,) * (len(sizes) - self.ndim) + self.shape
+    target = tuple(current if requested == -1 else requested for requested, current in zip(sizes, source))
+    return mx.broadcast_to(self.reshape(source), target)
+
+
+def _torch_any(self, dim=None, keepdim=False):
+    return _any(self, axis=dim, keepdims=keepdim)
+
+
+def _contiguous(self, memory_format=None):
+    return self
+
+
+def _requires_grad(self, requires_grad=True):
+    if requires_grad:
+        raise RuntimeError(
+            "requires_grad_ is not supported by the MLX backend; use torchmlx.Trainer"
+        )
+    return self
+
+
+def _backward(self, *args, **kwargs):
+    if args or kwargs:
+        raise TypeError("MLX backward compatibility does not accept arguments")
+    from ._autograd import backward
+
+    backward(self)
+
+
+def _torch_item(self):
+    from ._autograd import replayed_value
+
+    return _item(replayed_value(self))
+
+
+def _mask_indices(mask):
+    flat = mask.reshape(-1)
+    count = int(mx.sum(flat).item())
+    order = mx.argsort(flat.astype(mx.int32))
+    if count == 0:
+        return order[:0].astype(mx.int64)
+    return order[-count:].astype(mx.int64)
+
+
+def _torch_getitem(self, key):
+    if isinstance(key, mx.array) and key.dtype == mx.bool_:
+        indices = _mask_indices(key)
+        if key.shape == self.shape:
+            result = _getitem(self.reshape(-1), indices)
+            operation = lambda value: _getitem(_reshape(value, (-1,)), indices)
+        else:
+            result = _getitem(self, indices)
+            operation = lambda value: _getitem(value, indices)
+    else:
+        result = _getitem(self, key)
+        operation = lambda value: _getitem(value, key)
+    from ._autograd import propagate
+
+    return propagate(self, result, operation)
+
+
+def _torch_setitem(self, key, value):
+    if isinstance(key, mx.array) and key.dtype == mx.bool_:
+        indices = _mask_indices(key)
+        if key.shape == self.shape:
+            flat = self.reshape(-1)
+            _setitem(flat, indices, value)
+            return
+        _setitem(self, indices, value)
+        return
+    _setitem(self, key, value)
+
+
+def install(device_type):
+    mx.array.transpose = _torch_transpose
+    mx.array.reshape = _torch_reshape
+    mx.array.view = _view
+    mx.array.unsqueeze = _unsqueeze
+    mx.array.squeeze = _torch_squeeze
+    mx.array.flatten = _flatten
+    mx.array.float = _float
+    mx.array.bool = _bool
+    mx.array.pow = _pow
+    mx.array.mean = _torch_mean
+    mx.array.var = _torch_var
+    mx.array.size = _size
+    mx.array.chunk = _chunk
+    mx.array.to = _to
+    mx.array.repeat_interleave = _repeat_interleave
+    mx.array.masked_fill = _masked_fill
+    mx.array.expand = _expand
+    mx.array.any = _torch_any
+    mx.array.contiguous = _contiguous
+    mx.array.requires_grad_ = _requires_grad
+    mx.array.backward = _backward
+    mx.array.item = _torch_item
+    mx.array.device = property(lambda self: device_type("mps"))
+    mx.array.__getitem__ = _torch_getitem
+    mx.array.__setitem__ = _torch_setitem

torchmlx/nn/__init__.py

@@ -0,0 +1,257 @@
+from torchmlx._backend import BACKEND, unsupported
+
+
+if BACKEND == "torch":
+    import torch.nn as _native
+
+    Module = _native.Module
+    ModuleList = _native.ModuleList
+    Linear = _native.Linear
+    Embedding = _native.Embedding
+    LayerNorm = _native.LayerNorm
+    Sequential = _native.Sequential
+    GELU = _native.GELU
+    Dropout = _native.Dropout
+    Parameter = _native.Parameter
+
+    def __getattr__(name):
+        return getattr(_native, name)
+
+else:
+    from collections.abc import Iterable, MutableSequence
+
+    import mlx.core as mx
+    import mlx.nn as _native
+
+    class _ParameterTree(dict):
+        def __init__(self, values, model):
+            super().__init__(values)
+            self.model = model
+
+    class Module(_native.Module):
+        def __call__(self, *args, **kwargs):
+            from torchmlx._autograd import abort_forward, begin_forward, end_forward
+
+            begin_forward()
+            try:
+                output = self.forward(*args, **kwargs)
+            except Exception:
+                abort_forward()
+                raise
+            end_forward(self, args, kwargs, output)
+            return output
+
+        def forward(self, *args, **kwargs):
+            raise NotImplementedError(
+                f"Module [{type(self).__name__}] is missing the required forward function"
+            )
+
+        def parameters(self):
+            return _ParameterTree(super().parameters(), self)
+
+        def to(self, *args, **kwargs):
+            dtype = kwargs.pop("dtype", None)
+            device = kwargs.pop("device", None)
+            if kwargs:
+                name = next(iter(kwargs))
+                raise TypeError(f"to() got an unexpected keyword argument {name!r}")
+            for value in args:
+                if isinstance(value, mx.Dtype):
+                    if dtype is not None:
+                        raise TypeError("to() received dtype more than once")
+                    dtype = value
+                else:
+                    if device is not None:
+                        raise TypeError("to() received device more than once")
+                    device = value
+            if device is not None and str(device) != "mps":
+                raise ValueError("the MLX backend only accepts device='mps'")
+            if dtype is not None:
+                self.set_dtype(dtype)
+            return self
+
+        def register_buffer(self, name, tensor, persistent=True):
+            if not isinstance(name, str) or "." in name or name == "":
+                raise KeyError("buffer name must be a non-empty string without dots")
+            setattr(self, name, tensor)
+            self.freeze(recurse=False, keys=name, strict=True)
+
+    def Parameter(data=None, requires_grad=True):
+        if data is None:
+            return mx.array([])
+        if not requires_grad:
+            unsupported("torchmlx.nn.Parameter with requires_grad=False")
+        return data
+
+    class Linear(Module, _native.Linear):
+        def __init__(self, in_features, out_features, bias=True, device=None, dtype=None):
+            Module.__init__(self)
+            if device is not None and str(device) != "mps":
+                raise ValueError("the MLX backend only accepts device='mps'")
+            scale = (1 / in_features) ** 0.5
+            self.weight = mx.random.uniform(
+                low=-scale, high=scale, shape=(out_features, in_features)
+            )
+            if bias:
+                self.bias = mx.random.uniform(
+                    low=-scale, high=scale, shape=(out_features,)
+                )
+            if dtype is not None:
+                self.set_dtype(dtype)
+
+        def forward(self, input):
+            return _native.Linear.__call__(self, input)
+
+    class Embedding(Module, _native.Embedding):
+        def __init__(
+            self,
+            num_embeddings,
+            embedding_dim,
+            padding_idx=None,
+            max_norm=None,
+            norm_type=2.0,
+            scale_grad_by_freq=False,
+            sparse=False,
+            device=None,
+            dtype=None,
+        ):
+            if any(
+                value is not None
+                for value in (padding_idx, max_norm)
+            ) or norm_type != 2.0 or scale_grad_by_freq or sparse:
+                unsupported("torchmlx.nn.Embedding with non-default options")
+            Module.__init__(self)
+            if device is not None and str(device) != "mps":
+                raise ValueError("the MLX backend only accepts device='mps'")
+            scale = (1 / embedding_dim) ** 0.5
+            self.weight = mx.random.normal(
+                shape=(num_embeddings, embedding_dim), scale=scale
+            )
+            if dtype is not None:
+                self.set_dtype(dtype)
+
+        def forward(self, input):
+            return self.weight[input]
+
+    class LayerNorm(Module, _native.LayerNorm):
+        def __init__(
+            self,
+            normalized_shape,
+            eps=1e-5,
+            elementwise_affine=True,
+            bias=True,
+            device=None,
+            dtype=None,
+        ):
+            if isinstance(normalized_shape, Iterable) and not isinstance(
+                normalized_shape, (str, bytes)
+            ):
+                shape = tuple(normalized_shape)
+                if len(shape) != 1:
+                    unsupported("torchmlx.nn.LayerNorm with multidimensional shape")
+                normalized_shape = shape[0]
+            if device is not None and str(device) != "mps":
+                raise ValueError("the MLX backend only accepts device='mps'")
+            _native.LayerNorm.__init__(
+                self,
+                normalized_shape,
+                eps=eps,
+                affine=elementwise_affine,
+                bias=bias,
+            )
+            if dtype is not None:
+                self.set_dtype(dtype)
+
+        def forward(self, input):
+            return _native.LayerNorm.__call__(self, input)
+
+    class Sequential(Module):
+        def __init__(self, *args):
+            super().__init__()
+            self.layers = list(args)
+
+        def forward(self, input):
+            for layer in self.layers:
+                input = layer(input)
+            return input
+
+        def __len__(self):
+            return len(self.layers)
+
+        def __getitem__(self, index):
+            if isinstance(index, str):
+                return dict.__getitem__(self, index)
+            return self.layers[index]
+
+    class GELU(Module, _native.GELU):
+        def __init__(self, approximate="none"):
+            if approximate not in {"none", "tanh"}:
+                raise ValueError("approximate must be 'none' or 'tanh'")
+            _native.GELU.__init__(self, approx=approximate)
+
+        def forward(self, input):
+            return _native.GELU.__call__(self, input)
+
+    class Dropout(Module, _native.Dropout):
+        def __init__(self, p=0.5, inplace=False):
+            if inplace:
+                unsupported("torchmlx.nn.Dropout with inplace=True")
+            _native.Dropout.__init__(self, p=p)
+
+        def forward(self, input):
+            return _native.Dropout.__call__(self, input)
+
+    class ModuleList(Module, MutableSequence):
+        def __init__(self, modules=None):
+            super().__init__()
+            self.layers = list(modules or [])
+
+        def __getitem__(self, index):
+            if isinstance(index, str):
+                return dict.__getitem__(self, index)
+            return self.layers[index]
+
+        def __setitem__(self, index, module):
+            if isinstance(index, str):
+                dict.__setitem__(self, index, module)
+                return
+            self.layers[index] = module
+
+        def __delitem__(self, index):
+            if isinstance(index, str):
+                dict.__delitem__(self, index)
+                return
+            del self.layers[index]
+
+        def __len__(self):
+            return len(self.layers)
+
+        def __iter__(self):
+            return iter(self.layers)
+
+        def insert(self, index, module):
+            self.layers.insert(index, module)
+
+        def forward(self, *args, **kwargs):
+            unsupported("torchmlx.nn.ModuleList.forward")
+
+    def __getattr__(name):
+        unsupported(f"torchmlx.nn.{name}")
+
+
+import importlib as _importlib
+
+functional = _importlib.import_module("torchmlx.nn.functional")
+
+__all__ = [
+    "Embedding",
+    "Dropout",
+    "GELU",
+    "LayerNorm",
+    "Linear",
+    "Module",
+    "ModuleList",
+    "Parameter",
+    "Sequential",
+    "functional",
+]

torchmlx/nn/functional.py

@@ -0,0 +1,129 @@
+import math
+
+from torchmlx._backend import BACKEND, unsupported
+
+
+if BACKEND == "torch":
+    import torch.nn.functional as _native
+
+    cross_entropy = _native.cross_entropy
+    scaled_dot_product_attention = _native.scaled_dot_product_attention
+    silu = _native.silu
+    softmax = _native.softmax
+
+    def __getattr__(name):
+        return getattr(_native, name)
+
+else:
+    import mlx.core as mx
+
+    def cross_entropy(
+        input,
+        target,
+        weight=None,
+        size_average=None,
+        ignore_index=-100,
+        reduce=None,
+        reduction="mean",
+        label_smoothing=0.0,
+    ):
+        original_input = input
+
+        def finish(loss):
+            from torchmlx._autograd import register_loss
+
+            def rebuild(recomputed_input):
+                return cross_entropy(
+                    recomputed_input,
+                    target,
+                    weight=weight,
+                    size_average=size_average,
+                    ignore_index=ignore_index,
+                    reduce=reduce,
+                    reduction=reduction,
+                    label_smoothing=label_smoothing,
+                )
+
+            return register_loss(loss, original_input, rebuild)
+
+        if size_average is not None or reduce is not None:
+            unsupported("torchmlx.nn.functional.cross_entropy legacy reductions")
+        if input.ndim < 2:
+            raise ValueError("cross_entropy input must have at least 2 dimensions")
+        if input.ndim == 2:
+            logits = input
+            targets = target.reshape(-1)
+            output_shape = target.shape
+        else:
+            axes = [0] + list(range(2, input.ndim)) + [1]
+            logits = input.transpose(axes).reshape(-1, input.shape[1])
+            targets = target.reshape(-1)
+            output_shape = target.shape
+        valid = targets != ignore_index
+        safe_targets = mx.where(valid, targets, mx.zeros_like(targets))
+        log_probabilities = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
+        target_losses = -mx.take_along_axis(
+            log_probabilities, safe_targets[:, None], axis=-1
+        ).squeeze(-1)
+        if weight is None:
+            smooth_losses = -mx.mean(log_probabilities, axis=-1)
+        else:
+            target_losses = target_losses * weight[safe_targets]
+            smooth_losses = -mx.sum(
+                log_probabilities * weight[None, :], axis=-1
+            ) / logits.shape[-1]
+        losses = (1 - label_smoothing) * target_losses + label_smoothing * smooth_losses
+        losses = mx.where(valid, losses, mx.zeros_like(losses))
+        if reduction == "none":
+            return finish(losses.reshape(output_shape))
+        if reduction == "sum":
+            return finish(mx.sum(losses))
+        if reduction == "mean":
+            if weight is None:
+                denominator = mx.sum(valid)
+            else:
+                denominator = mx.sum(mx.where(valid, weight[safe_targets], 0))
+            return finish(mx.sum(losses) / denominator)
+        raise ValueError(f"invalid reduction {reduction!r}")
+
+    def scaled_dot_product_attention(
+        query,
+        key,
+        value,
+        attn_mask=None,
+        dropout_p=0.0,
+        is_causal=False,
+        scale=None,
+        enable_gqa=False,
+    ):
+        if dropout_p != 0.0:
+            unsupported(
+                "torchmlx.nn.functional.scaled_dot_product_attention with dropout"
+            )
+        if enable_gqa:
+            unsupported(
+                "torchmlx.nn.functional.scaled_dot_product_attention with enable_gqa"
+            )
+        if is_causal and attn_mask is not None:
+            raise ValueError("attn_mask cannot be set when is_causal=True")
+        if scale is None:
+            scale = 1 / math.sqrt(query.shape[-1])
+        mask = "causal" if is_causal else attn_mask
+        return mx.fast.scaled_dot_product_attention(
+            query, key, value, scale=scale, mask=mask
+        )
+
+    def silu(input, inplace=False):
+        if inplace:
+            unsupported("torchmlx.nn.functional.silu with inplace=True")
+        return input * mx.sigmoid(input)
+
+    def softmax(input, dim=None, dtype=None):
+        value = input.astype(dtype) if dtype is not None else input
+        return mx.softmax(value, axis=dim)
+
+    def __getattr__(name):
+        unsupported(f"torchmlx.nn.functional.{name}")
+
+
+__all__ = ["cross_entropy", "scaled_dot_product_attention", "silu", "softmax"]

torchmlx/optim/__init__.py

@@ -0,0 +1,72 @@
+from torchmlx._backend import BACKEND, unsupported
+
+
+if BACKEND == "torch":
+    import torch.optim as _native
+
+    AdamW = _native.AdamW
+
+    def __getattr__(name):
+        return getattr(_native, name)
+
+else:
+    import mlx.optimizers as _native
+
+    class AdamW:
+        def __init__(
+            self,
+            params,
+            lr=1e-3,
+            betas=(0.9, 0.999),
+            eps=1e-8,
+            weight_decay=1e-2,
+            amsgrad=False,
+            maximize=False,
+            foreach=None,
+            capturable=False,
+            differentiable=False,
+            fused=None,
+        ):
+            if amsgrad or maximize or foreach is not None or capturable or differentiable or fused is not None:
+                unsupported("torchmlx.optim.AdamW with non-default options")
+            model = getattr(params, "model", None)
+            if model is None:
+                raise TypeError(
+                    "MLX AdamW requires parameters returned directly by model.parameters()"
+                )
+            self._parameters = params
+            self._model = model
+            self._optimizer = _native.AdamW(
+                learning_rate=lr,
+                betas=list(betas),
+                eps=eps,
+                weight_decay=weight_decay,
+                bias_correction=True,
+            )
+            self._pending_update = None
+            self._random_before_forward = None
+
+        @property
+        def state(self):
+            return self._optimizer.state
+
+        def zero_grad(self, *args, **kwargs):
+            if args or kwargs:
+                unsupported("torchmlx.optim.AdamW.zero_grad with arguments")
+            from torchmlx._autograd import activate
+
+            self._pending_update = None
+            activate(self)
+
+        def step(self, *args, **kwargs):
+            if args or kwargs:
+                unsupported("torchmlx.optim.AdamW.step with arguments")
+            from torchmlx._autograd import step
+
+            step(self)
+
+    def __getattr__(name):
+        unsupported(f"torchmlx.optim.{name}")
+
+
+__all__ = ["AdamW"]

torchmlx/trainer.py

@@ -0,0 +1,60 @@
+from ._backend import BACKEND
+
+
+if BACKEND == "torch":
+    import torch
+
+    class Trainer:
+        def __init__(self, model, optimizer, loss_fn, compile=True):
+            self.model = model
+            self.optimizer = optimizer
+            self.loss_fn = loss_fn
+            self._step = torch.compile(self._train_step) if compile else self._train_step
+
+        def _train_step(self, x, y):
+            self.optimizer.zero_grad()
+            loss = self.loss_fn(self.model(x), y)
+            loss.backward()
+            self.optimizer.step()
+            return loss
+
+        def step(self, x, y):
+            return self._step(x, y)
+
+else:
+    import mlx.core as mx
+    import mlx.nn as nn
+
+    class Trainer:
+        def __init__(self, model, optimizer, loss_fn, compile=True):
+            self.model = model
+            self.optimizer = optimizer
+            self.loss_fn = loss_fn
+            native_optimizer = optimizer._optimizer
+            native_optimizer.init(model.trainable_parameters())
+
+            def objective(x, y):
+                return loss_fn(model(x), y)
+
+            value_and_grad = nn.value_and_grad(model, objective)
+
+            def train_step(x, y):
+                loss, gradients = value_and_grad(x, y)
+                native_optimizer.update(model, gradients)
+                return loss
+
+            if compile:
+                state = [model.state, native_optimizer.state, mx.random.state]
+                self._step = mx.compile(train_step, inputs=state, outputs=state)
+            else:
+                self._step = train_step
+
+        def step(self, x, y):
+            loss = self._step(x, y)
+            mx.eval(
+                loss,
+                self.model.parameters(),
+                self.optimizer._optimizer.state,
+                mx.random.state,
+            )
+            return loss