tensorless 0.4.0__tar.gz → 0.6.0__tar.gz

{tensorless-0.4.0 → tensorless-0.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tensorless
-Version: 0.4.0
+Version: 0.6.0
 Summary: ML with maximum automation and minimum setup.
 Author: Tensorless Contributors
 License: MIT
@@ -10,11 +10,25 @@ License-File: LICENSE
 Requires-Dist: numpy>=1.24
 Provides-Extra: dev
 Requires-Dist: pytest>=7.0; extra == "dev"
+Provides-Extra: jax
+Requires-Dist: jax>=0.4.30; extra == "jax"
+Provides-Extra: cuda
+Requires-Dist: jax[cuda12]>=0.4.30; extra == "cuda"
+Provides-Extra: tpu
+Requires-Dist: jax[tpu]>=0.4.30; extra == "tpu"
+Provides-Extra: mps
+Requires-Dist: mlx>=0.18; extra == "mps"
+Provides-Extra: accelerators
+Requires-Dist: jax[cuda12]>=0.4.30; extra == "accelerators"
+Requires-Dist: jax[tpu]>=0.4.30; extra == "accelerators"
+Requires-Dist: mlx>=0.18; extra == "accelerators"
 Dynamic: license-file
 
 # Tensorless
 
-Tensorless trains small native NumPy models with sensible defaults. It supports
+Tensorless trains small custom models with sensible defaults. It uses a native
+NumPy engine on CPU and optional JAX or MLX backends for accelerator execution.
+It supports
 text generation, text classification, tabular classification, and regression.
 
 ## Install
@@ -23,6 +37,17 @@ text generation, text classification, tabular classification, and regression.
 pip install -e .
 ```
 
+Optional accelerator backends:
+
+```bash
+pip install -e '.[cuda]'   # JAX CUDA
+pip install -e '.[tpu]'    # JAX TPU
+pip install -e '.[mps]'    # Apple Silicon MLX
+```
+
+CUDA and TPU backends currently accelerate transformer text tasks. Tabular
+tasks and unsupported platforms use the native CPU engine.
+
 ## Train on your data
 
 ```python
@@ -73,5 +98,11 @@ model = tl.load("model.tl")
 print(model.info())
 ```
 
+The native extension API is in `tensorless.engine`: `Module`, `Parameter`,
+`Adam`, and `SGD` provide model parameters, gradients, and optimization
+without a PyTorch dependency. Accelerator cache helpers are available as
+`tensorless.devices.clear_memory()` and `tensorless.devices.memory_stats()`.
+
 See the [documentation](docs/quickstart.md) for data formats, configuration,
 checkpointing, and the command-line interface.
+pypi-AgEIcHlwaS5vcmcCJDJmYzZiYmQ1LTI1YTAtNGNlZi05OWE2LTliMjg3ZjY5MThiZQACElsxLFsidGVuc29ybGVzcyJdXQACLFsyLFsiNjcxZGZmNDQtZmVmMC00MWNiLWIzYzYtZWQzMTI3OWU4NWM5Il1dAAAGIDZAp3AAlp0idrTOMPJ227qF_7W0LDsyUpi6pfqf0aV0

tensorless-0.4.0/tensorless.egg-info/PKG-INFO → tensorless-0.6.0/README.md

@@ -1,20 +1,8 @@
-Metadata-Version: 2.4
-Name: tensorless
-Version: 0.4.0
-Summary: ML with maximum automation and minimum setup.
-Author: Tensorless Contributors
-License: MIT
-Requires-Python: >=3.9
-Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: numpy>=1.24
-Provides-Extra: dev
-Requires-Dist: pytest>=7.0; extra == "dev"
-Dynamic: license-file
-
 # Tensorless
 
-Tensorless trains small native NumPy models with sensible defaults. It supports
+Tensorless trains small custom models with sensible defaults. It uses a native
+NumPy engine on CPU and optional JAX or MLX backends for accelerator execution.
+It supports
 text generation, text classification, tabular classification, and regression.
 
 ## Install
@@ -23,6 +11,17 @@ text generation, text classification, tabular classification, and regression.
 pip install -e .
 ```
 
+Optional accelerator backends:
+
+```bash
+pip install -e '.[cuda]'   # JAX CUDA
+pip install -e '.[tpu]'    # JAX TPU
+pip install -e '.[mps]'    # Apple Silicon MLX
+```
+
+CUDA and TPU backends currently accelerate transformer text tasks. Tabular
+tasks and unsupported platforms use the native CPU engine.
+
 ## Train on your data
 
 ```python
@@ -73,5 +72,11 @@ model = tl.load("model.tl")
 print(model.info())
 ```
 
+The native extension API is in `tensorless.engine`: `Module`, `Parameter`,
+`Adam`, and `SGD` provide model parameters, gradients, and optimization
+without a PyTorch dependency. Accelerator cache helpers are available as
+`tensorless.devices.clear_memory()` and `tensorless.devices.memory_stats()`.
+
 See the [documentation](docs/quickstart.md) for data formats, configuration,
 checkpointing, and the command-line interface.
+pypi-AgEIcHlwaS5vcmcCJDJmYzZiYmQ1LTI1YTAtNGNlZi05OWE2LTliMjg3ZjY5MThiZQACElsxLFsidGVuc29ybGVzcyJdXQACLFsyLFsiNjcxZGZmNDQtZmVmMC00MWNiLWIzYzYtZWQzMTI3OWU4NWM5Il1dAAAGIDZAp3AAlp0idrTOMPJ227qF_7W0LDsyUpi6pfqf0aV0

{tensorless-0.4.0 → tensorless-0.6.0}/docs/architecture.md

@@ -36,8 +36,12 @@ tensorless/
 │   └── manager.py                      CheckpointManager: atomic save/load/clear
 ├── serialization/
 │   └── tl_format.py                     save_tl/load_tl: the .tl file format
+├── backends/
+│   ├── jax_backend.py                    optional CUDA/TPU transformer backend and local data parallelism
+│   └── mlx_backend.py                    optional Apple Silicon MPS backend
 ├── devices/
-│   └── device.py                         native CPU device resolution
+│   ├── device.py                         accelerator detection and resolution
+│   └── memory.py                         backend-neutral cache and memory controls
 └── cli/
     └── main.py                            argparse-based CLI
 ```
@@ -60,7 +64,7 @@ auto.config.resolve_config(ds, TrainConfig) -> ResolvedConfig (every field concr
 training.data_prep.prepare_*(ds, cfg)   -> PreparedData (train/val DataLoaders,
   │                                          meta, tokenizer/preprocessor)
   ▼
-models.registry.build_model(task, model_type, cfg, meta) -> nn.Module
+  models.registry.build_model(task, model_type, cfg, meta) -> native/backend model
   │
   ▼
 training.trainer.run_training(...)      -> trains, checkpoints periodically,
@@ -97,6 +101,9 @@ matching the dataset's fingerprint.
 - **A `.tl` file is the unit of portability.** Nothing about inference
   should require the original dataset, training script, or checkpoint
   directory to still exist.
+- **The public model API is backend-neutral.** `engine.Module` and
+  `engine.Parameter` expose parameters and gradients consistently while JAX
+  and MLX provide accelerator execution behind the same model contract.
 
 ## Extending Tensorless
 
@@ -120,7 +127,16 @@ for the new extension, producing a `Dataset` with the appropriate `kind`.
 ### Adding a new backend/device
 
 Extend `devices/device.py`'s `_*_available()` checks and
-`auto_select_device()` / `get_torch_device()`.
+`auto_select_device()` / `get_device()`. Add a lazy-imported backend module
+and select it in `models/registry.py`, `training/trainer.py`, and
+`runtime.py`. Backend models must expose the native `state_dict()` contract so
+`.tl` files remain portable.
+
+When multiple local JAX CUDA or TPU devices are visible, the JAX transformer
+splits each compatible batch with `pmap` and averages gradients with a device
+collective. For multi-host jobs, set `JAX_COORDINATOR_ADDRESS`,
+`JAX_PROCESS_COUNT`, and `JAX_PROCESS_ID` consistently across processes before
+launching. Single-device training uses the same model without collectives.
 
 See [contributing.md](contributing.md) for the contribution process
 itself (tests, PRs, etc).

{tensorless-0.4.0 → tensorless-0.6.0}/docs/checkpointing.md

@@ -12,9 +12,11 @@ create or manage this directory yourself.
 `checkpoint.pt`, containing:
 
 - `model_state_dict` — model weights
+- `best_model_state_dict` — best validation weights, when validation is enabled
 - `optimizer_state_dict` — optimizer momentum/variance buffers
 - `scheduler_state_dict` — learning rate schedule position
 - `epoch`, `global_step` — where training left off
+- `train_loader_epoch` — deterministic shuffle position for resumed training
 - `early_stopping_best`, `early_stopping_bad_checks` — early stopping state
 - `config` — the fully resolved training configuration used
 - `meta` — task-specific sizing info (vocab size, number of classes, etc.)
@@ -28,6 +30,12 @@ final `.tl` file — nothing about resumption depends on the original
 dataset still being on disk in the same location, only on it being
 fingerprint-identical to what was originally used.
 
+Set `gradient_checkpointing=True` to reduce native transformer activation
+memory. Backward recomputes each block and replays its dropout state, trading
+extra compute for lower peak memory. The JAX backend uses `jax.checkpoint`
+for the same option and uses device collectives for local multi-device
+gradient reduction.
+
 ## When checkpoints are written
 
 - Every `checkpoint_every` steps (default: 50) during training, with
@@ -40,6 +48,13 @@ Writes are atomic: Tensorless writes to a temporary file in the same
 directory and renames it into place, so a crash mid-write never leaves a
 corrupt checkpoint that would block resumption.
 
+JAX and MLX accelerator arrays are converted to portable NumPy arrays before
+they enter a `.tl` file or checkpoint. This keeps models loadable on a CPU
+machine. Set `checkpoint_shard_size_mb` above zero to store model and
+optimizer state in independently readable shard files plus a small
+`checkpoint.pt` manifest. The normal resume API reconstructs the state
+automatically.
+
 When loading `.tl` files, Tensorless fills compatible fields introduced by
 older versions with safe defaults. Files created by a newer unsupported format
 version are rejected with an upgrade message instead of being partially read.

{tensorless-0.4.0 → tensorless-0.6.0}/docs/configuration.md

@@ -25,6 +25,7 @@ defaults.
 | `heads` | auto | Attention heads (transformer only) |
 | `ff_mult` | auto | Feed-forward expansion multiplier (transformer only) |
 | `dropout` | `0.1` | Dropout probability |
+| `gradient_checkpointing` | `False` | Recompute transformer activations during backward to reduce memory use; native and JAX paths support this |
 | `max_seq_len` | `256` (text) / `1` (tabular) | Max sequence length in tokens (text tasks) |
 | `tokenizer` | `"bpe"` | Text tokenizer: `"bpe"` or `"char"` |
 | `bpe_vocab_size` | `1000` | Maximum vocabulary size when `tokenizer="bpe"` |
@@ -54,8 +55,8 @@ defaults.
 
 | Field | Default | Description |
 |---|---|---|
-| `device` | auto (`tpu` > `cuda` > `mps` > `cpu`) | Force a specific device |
-| `precision` | auto | `"fp32"`, `"fp16"`, or `"bf16"` |
+| `device` | auto (`tpu` > `cuda` > `mps` > `cpu`) | Force a specific device; accelerator backends require their optional package |
+| `precision` | auto | `"fp32"`, `"fp16"`, or `"bf16"`; JAX and MLX accelerator transformers use the selected dtype, fp32 loss math, and fp16 loss scaling |
 
 ## Checkpointing
 
@@ -63,6 +64,7 @@ defaults.
 |---|---|---|
 | `checkpoint_every` | `50` | Steps between checkpoint writes |
 | `checkpoint_dir` | `"<out>.ckpt"` | Checkpoint directory |
+| `checkpoint_shard_size_mb` | `0` | Split model and optimizer checkpoint state into parts of approximately this size; `0` keeps the single-file format |
 
 ## Misc
 
@@ -93,3 +95,8 @@ tl.train(
 Passing an unrecognized keyword raises a `ConfigError` listing every
 valid field name, so typos are caught immediately rather than silently
 ignored.
+
+For backend-neutral accelerator maintenance, use
+`tensorless.devices.clear_memory()` and
+`tensorless.devices.memory_stats()`. These helpers use JAX or MLX when
+available and otherwise safely do nothing.

{tensorless-0.4.0 → tensorless-0.6.0}/docs/roadmap.md

@@ -5,8 +5,9 @@ a commitment or timeline.
 
 ## Near-term
 
-- **Progress bars** for training (currently plain print-based logging)
-- **Multi-GPU / distributed training** for larger datasets
+- **Hyperparameter search mode**: an opt-in `tl.train("./data",
+  search=True)` that tries a small set of configurations and keeps the
+  best, rather than a single heuristic choice
 
 ## Medium-term
 
@@ -15,18 +16,20 @@ a commitment or timeline.
   from scratch
 - **Additional data formats**: Parquet, Excel (`.xlsx`), image
   directories, audio
-- **Hyperparameter search mode**: an opt-in `tl.train("./data",
-  search=True)` that tries a small set of configurations and keeps the
-  best, rather than a single heuristic choice
 - **Data quality auto-fixes**: currently `tl.inspect()` only *reports*
   problems like missing values or class imbalance; a future mode could
   offer to fix them (with explicit user opt-in, consistent with "never
   silently modify user data")
 
-## Long-term / exploratory
+## Implemented
 
-- **Alternate backends** (JAX, a lightweight NumPy-only backend) behind
-  the same `tl.train()`/`tl.load()` API
+- **Progress bars** for batch and epoch training output
+- **JAX CUDA/TPU backend**, including local and multi-host data parallelism
+- **MLX Apple Silicon backend** for transformer text tasks
+- **NumPy CPU backend** with stacked attention, dropout, gradient checkpointing,
+  and sharded checkpoints
+
+## Long-term / exploratory
 - **Export to other formats** (ONNX, TorchScript) from a `.tl` file for
   deployment outside Python
 

{tensorless-0.4.0 → tensorless-0.6.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tensorless"
-version = "0.4.0"
+version = "0.6.0"
 description = "ML with maximum automation and minimum setup."
 readme = "README.md"
 requires-python = ">=3.9"
@@ -16,6 +16,11 @@ dependencies = [
 
 [project.optional-dependencies]
 dev = ["pytest>=7.0"]
+jax = ["jax>=0.4.30"]
+cuda = ["jax[cuda12]>=0.4.30"]
+tpu = ["jax[tpu]>=0.4.30"]
+mps = ["mlx>=0.18"]
+accelerators = ["jax[cuda12]>=0.4.30", "jax[tpu]>=0.4.30", "mlx>=0.18"]
 
 [project.scripts]
 tensorless = "tensorless.cli.main:main"

{tensorless-0.4.0 → tensorless-0.6.0}/tensorless/api.py

@@ -179,6 +179,7 @@ def pretrain(
 
 
 def _finalize_from_checkpoint(ckpt: dict, out: str, verbose: bool) -> LoadedModel:
+    model_state = ckpt.get("best_model_state_dict") or ckpt["model_state_dict"]
     payload = {
         "tl_format_version": _tl_format_version,
         "tensorless_version": _tl_version,
@@ -186,7 +187,7 @@ def _finalize_from_checkpoint(ckpt: dict, out: str, verbose: bool) -> LoadedMode
         "model_type": ckpt["config"]["model_type"],
         "config": ckpt["config"],
         "meta": ckpt["meta"],
-        "model_state_dict": ckpt["model_state_dict"],
+        "model_state_dict": model_state,
         "tokenizer_state": ckpt.get("tokenizer_state"),
         "preprocessor_state": ckpt.get("preprocessor_state"),
         "dataset_fingerprint": ckpt["dataset_fingerprint"],

{tensorless-0.4.0 → tensorless-0.6.0}/tensorless/auto/config.py

@@ -111,6 +111,7 @@ def resolve_config(ds: Dataset, user: TrainConfig) -> ResolvedConfig:
         heads=user.heads or heads,
         ff_mult=user.ff_mult or ff_mult,
         dropout=user.dropout if user.dropout is not None else 0.1,
+        gradient_checkpointing=bool(user.gradient_checkpointing),
         max_seq_len=max_seq_len,
         tokenizer=tokenizer,
         bpe_vocab_size=user.bpe_vocab_size if user.bpe_vocab_size is not None else _auto_vocab_size(ds),
@@ -129,6 +130,7 @@ def resolve_config(ds: Dataset, user: TrainConfig) -> ResolvedConfig:
         precision=precision,
         checkpoint_every=user.checkpoint_every or 50,
         checkpoint_dir=checkpoint_dir,
+        checkpoint_shard_size_mb=user.checkpoint_shard_size_mb or 0,
         seed=user.seed,
         verbose=user.verbose,
     )

tensorless-0.6.0/tensorless/backends/__init__.py

@@ -0,0 +1 @@
+"""Optional execution backends."""

tensorless-0.6.0/tensorless/backends/jax_backend.py

@@ -0,0 +1,165 @@
+"""JAX implementation of the text transformer.
+
+JAX is imported lazily so importing Tensorless remains CPU/NumPy-only when the
+optional dependency is not installed.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+import os
+
+import numpy as np
+
+from ..models.transformer import TinyTransformer
+
+
+def _jax():
+    try:
+        import jax
+        import jax.numpy as jnp
+        distributed_vars = ("JAX_COORDINATOR_ADDRESS", "JAX_PROCESS_COUNT", "JAX_PROCESS_ID")
+        if all(name in os.environ for name in distributed_vars) and not jax.distributed.is_initialized():
+            jax.distributed.initialize(
+                coordinator_address=os.environ["JAX_COORDINATOR_ADDRESS"],
+                num_processes=int(os.environ["JAX_PROCESS_COUNT"]),
+                process_id=int(os.environ["JAX_PROCESS_ID"]),
+            )
+    except ImportError as exc:
+        raise ImportError("JAX is required for the cuda/tpu backend; install tensorless[jax].") from exc
+    return jax, jnp
+
+
+def is_available() -> bool:
+    try:
+        _jax()
+        return True
+    except ImportError:
+        return False
+
+
+class JaxTinyTransformer(TinyTransformer):
+    """The native transformer parameter layout with JAX math and gradients."""
+
+    def __init__(self, *args, **kwargs):
+        self.precision = kwargs.pop("precision", "fp32")
+        super().__init__(*args, **kwargs)
+        d_model = kwargs.get("d_model", args[1] if len(args) > 1 else None)
+        self.heads = kwargs.get("heads", args[3] if len(args) > 3 else None)
+        if d_model is None or self.heads is None:
+            raise ValueError("JAX transformer requires d_model and heads")
+        self.d_model = d_model
+        self.head_dim = d_model // self.heads
+
+    def _jax_params(self):
+        _, jnp = _jax()
+        dtype = {"fp16": jnp.float16, "bf16": jnp.bfloat16}.get(self.precision, jnp.float32)
+        return {name: jnp.asarray(value.data, dtype=dtype) for name, value in self.named_parameters()}
+
+    def _forward_jax(self, params, input_ids, attention_mask=None):
+        _, jnp = _jax()
+
+        def layer_norm(x, gain, bias, eps=1e-5):
+            mean = jnp.mean(x, axis=-1, keepdims=True)
+            var = jnp.var(x, axis=-1, keepdims=True)
+            return (x - mean) / jnp.sqrt(var + eps) * gain + bias
+
+        input_ids = jnp.asarray(input_ids, dtype=jnp.int32)
+        batch, length = input_ids.shape
+        hidden = params["tok_emb"][input_ids] + params["pos_emb"][jnp.arange(length)]
+        causal = jnp.tril(jnp.ones((length, length), dtype=bool))
+        for index in range(len(self.blocks)):
+            prefix = f"blocks.{index}"
+            normed1 = layer_norm(hidden, params[f"{prefix}.ln1_gain"], params[f"{prefix}.ln1_bias"])
+            qkv = normed1 @ params[f"{prefix}.qkv_weight"] + params[f"{prefix}.qkv_bias"]
+            q, k, v = jnp.split(qkv, 3, axis=-1)
+            q = q.reshape(batch, length, self.heads, self.head_dim).transpose(0, 2, 1, 3)
+            k = k.reshape(batch, length, self.heads, self.head_dim).transpose(0, 2, 1, 3)
+            v = v.reshape(batch, length, self.heads, self.head_dim).transpose(0, 2, 1, 3)
+            scores = q @ k.transpose(0, 1, 3, 2) / np.sqrt(self.head_dim)
+            scores = jnp.where(causal[None, None, :, :], scores, -1e9)
+            if attention_mask is not None:
+                mask = jnp.asarray(attention_mask) > 0
+                scores = jnp.where(mask[:, None, None, :], scores, -1e9)
+            probabilities = jnp.exp(scores - scores.max(axis=-1, keepdims=True))
+            probabilities /= probabilities.sum(axis=-1, keepdims=True)
+            context = probabilities @ v
+            context = context.transpose(0, 2, 1, 3).reshape(batch, length, self.d_model)
+            attention = context @ params[f"{prefix}.out_weight"] + params[f"{prefix}.out_bias"]
+            residual = hidden + attention
+            normed2 = layer_norm(residual, params[f"{prefix}.ln2_gain"], params[f"{prefix}.ln2_bias"])
+            ff_pre = normed2 @ params[f"{prefix}.ff1_weight"] + params[f"{prefix}.ff1_bias"]
+            ff_hidden = jnp.maximum(ff_pre, 0)
+            ff = ff_hidden @ params[f"{prefix}.ff2_weight"] + params[f"{prefix}.ff2_bias"]
+            hidden = residual + ff
+        hidden = layer_norm(hidden, params["ln_f_gain"], params["ln_f_bias"])
+        if self.task == "text-generation":
+            return hidden @ params["tok_emb"].T + params["head_bias"]
+        mask = jnp.ones((batch, length), dtype=jnp.float32) if attention_mask is None else jnp.asarray(attention_mask)
+        pooled = (hidden * mask[:, :, None]).sum(1) / jnp.maximum(mask.sum(1, keepdims=True), 1)
+        return pooled @ params["head_weight"] + params["head_bias"]
+
+    def forward(self, input_ids, attention_mask=None, cache=False):
+        _, jnp = _jax()
+        output = self._forward_jax(self._jax_params(), input_ids, attention_mask)
+        output = np.asarray(output)
+        return (output, None) if cache else output
+
+    def loss_and_backward(self, input_ids, target, attention_mask=None):
+        jax, jnp = _jax()
+        params = self._jax_params()
+        target = jnp.asarray(target, dtype=jnp.int32)
+
+        def loss_fn(current, current_ids, current_target, current_mask):
+            forward = jax.checkpoint(self._forward_jax) if self.gradient_checkpointing else self._forward_jax
+            logits = forward(current, current_ids, current_mask)
+            flat_logits = logits.reshape(-1, logits.shape[-1]).astype(jnp.float32)
+            flat_target = current_target.reshape(-1)
+            log_probs = jax.nn.log_softmax(flat_logits, axis=-1)
+            losses = -jnp.take_along_axis(log_probs, flat_target[:, None], axis=1).squeeze(1)
+            if self.task == "text-generation":
+                valid = flat_target != self.pad_id
+                return jnp.sum(jnp.where(valid, losses, 0.0)) / jnp.maximum(valid.sum(), 1)
+            return jnp.mean(losses)
+
+        loss_scale = 128.0 if self.precision == "fp16" else 1.0
+        scaled_loss_fn = lambda *arguments: loss_fn(*arguments) * loss_scale
+        mask = jnp.ones_like(input_ids, dtype=jnp.float32) if attention_mask is None else jnp.asarray(attention_mask)
+        devices = jax.local_device_count()
+        if devices > 1 and input_ids.shape[0] % devices == 0:
+            per_device_loss = jax.value_and_grad(scaled_loss_fn)
+
+            def mapped_step(current, ids, labels, current_mask):
+                value, gradients = per_device_loss(current, ids, labels, current_mask)
+                return jax.lax.pmean(value, "data"), jax.tree_util.tree_map(
+                    lambda gradient: jax.lax.pmean(gradient, "data"), gradients
+                )
+
+            shard = lambda value: value.reshape((devices, value.shape[0] // devices) + value.shape[1:])
+            loss, gradients = jax.pmap(mapped_step, axis_name="data", in_axes=(None, 0, 0, 0))(
+                params, shard(jnp.asarray(input_ids)), shard(target), shard(mask)
+            )
+            loss = loss[0]
+            gradients = jax.tree_util.tree_map(lambda gradient: gradient[0] / loss_scale, gradients)
+        else:
+            loss, gradients = jax.value_and_grad(scaled_loss_fn)(params, input_ids, target, mask)
+            gradients = jax.tree_util.tree_map(lambda gradient: gradient / loss_scale, gradients)
+        for name, parameter in self.named_parameters():
+            parameter.grad[...] = np.asarray(gradients[name], dtype=np.float32)
+        return float(loss) / loss_scale
+
+    def generate(self, input_ids, max_new_tokens, temperature=.8, top_k=40, eos_id: Optional[int] = None):
+        ids = np.asarray(input_ids, dtype=np.int64).copy()
+        for _ in range(max_new_tokens):
+            logits = self.forward(ids[:, -self.max_seq_len:])[:, -1, :] / max(temperature, 1e-5)
+            if top_k is not None:
+                k = min(top_k, logits.shape[-1])
+                excluded = np.argpartition(logits, -k, axis=1)[:, :-k]
+                logits[np.arange(len(ids))[:, None], excluded] = -np.inf
+            probs = np.exp(logits - logits.max(1, keepdims=True))
+            probs /= probs.sum(1, keepdims=True)
+            next_ids = np.array([np.random.choice(logits.shape[1], p=p) for p in probs])[:, None]
+            ids = np.concatenate([ids, next_ids], axis=1)
+            if eos_id is not None and np.all(next_ids == eos_id):
+                break
+        return ids

tensorless-0.6.0/tensorless/backends/mlx_backend.py

@@ -0,0 +1,133 @@
+"""MLX implementation of the text transformer.
+
+MLX is imported lazily so Tensorless remains usable on non-Apple systems
+without installing an accelerator framework.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import numpy as np
+
+from ..models.transformer import TinyTransformer
+
+
+def _mlx():
+    try:
+        import mlx.core as mx
+    except ImportError as exc:
+        raise ImportError("MLX is required for the mps backend; install tensorless[mps].") from exc
+    return mx
+
+
+def is_available() -> bool:
+    try:
+        _mlx()
+        return True
+    except ImportError:
+        return False
+
+
+class MlxTinyTransformer(TinyTransformer):
+    """The native Tensorless parameter layout with MLX math and gradients."""
+
+    def __init__(self, *args, **kwargs):
+        self.precision = kwargs.pop("precision", "fp32")
+        super().__init__(*args, **kwargs)
+        d_model = kwargs.get("d_model", args[1] if len(args) > 1 else None)
+        self.heads = kwargs.get("heads", args[3] if len(args) > 3 else None)
+        if d_model is None or self.heads is None:
+            raise ValueError("MLX transformer requires d_model and heads")
+        self.d_model = d_model
+        self.head_dim = d_model // self.heads
+
+    def _mlx_params(self):
+        mx = _mlx()
+        dtype = {"fp16": mx.float16, "bf16": mx.bfloat16}.get(self.precision, mx.float32)
+        return {name: mx.array(parameter.data, dtype=dtype) for name, parameter in self.named_parameters()}
+
+    def _forward_mlx(self, params, input_ids, attention_mask=None):
+        mx = _mlx()
+
+        def layer_norm(x, gain, bias, eps=1e-5):
+            mean = mx.mean(x, axis=-1, keepdims=True)
+            var = mx.var(x, axis=-1, keepdims=True)
+            return (x - mean) / mx.sqrt(var + eps) * gain + bias
+
+        input_ids = mx.array(input_ids, dtype=mx.int32)
+        batch, length = input_ids.shape
+        hidden = params["tok_emb"][input_ids] + params["pos_emb"][mx.arange(length)]
+        causal = mx.tril(mx.ones((length, length), dtype=mx.bool_))
+        for index in range(len(self.blocks)):
+            prefix = f"blocks.{index}"
+            normed1 = layer_norm(hidden, params[f"{prefix}.ln1_gain"], params[f"{prefix}.ln1_bias"])
+            qkv = normed1 @ params[f"{prefix}.qkv_weight"] + params[f"{prefix}.qkv_bias"]
+            q, k, v = mx.split(qkv, 3, axis=-1)
+            q = q.reshape(batch, length, self.heads, self.head_dim).transpose(0, 2, 1, 3)
+            k = k.reshape(batch, length, self.heads, self.head_dim).transpose(0, 2, 1, 3)
+            v = v.reshape(batch, length, self.heads, self.head_dim).transpose(0, 2, 1, 3)
+            scores = q @ k.transpose(0, 1, 3, 2) / np.sqrt(self.head_dim)
+            scores = mx.where(causal[None, None, :, :], scores, -1e9)
+            if attention_mask is not None:
+                mask = mx.array(attention_mask) > 0
+                scores = mx.where(mask[:, None, None, :], scores, -1e9)
+            probabilities = mx.softmax(scores, axis=-1)
+            context = probabilities @ v
+            context = context.transpose(0, 2, 1, 3).reshape(batch, length, self.d_model)
+            attention = context @ params[f"{prefix}.out_weight"] + params[f"{prefix}.out_bias"]
+            residual = hidden + attention
+            normed2 = layer_norm(residual, params[f"{prefix}.ln2_gain"], params[f"{prefix}.ln2_bias"])
+            ff_pre = normed2 @ params[f"{prefix}.ff1_weight"] + params[f"{prefix}.ff1_bias"]
+            ff_hidden = mx.maximum(ff_pre, 0)
+            ff = ff_hidden @ params[f"{prefix}.ff2_weight"] + params[f"{prefix}.ff2_bias"]
+            hidden = residual + ff
+        hidden = layer_norm(hidden, params["ln_f_gain"], params["ln_f_bias"])
+        if self.task == "text-generation":
+            return hidden @ params["tok_emb"].T + params["head_bias"]
+        mask = mx.ones((batch, length)) if attention_mask is None else mx.array(attention_mask)
+        pooled = (hidden * mask[:, :, None]).sum(1) / mx.maximum(mask.sum(1, keepdims=True), 1)
+        return pooled @ params["head_weight"] + params["head_bias"]
+
+    def forward(self, input_ids, attention_mask=None, cache=False):
+        output = np.asarray(self._forward_mlx(self._mlx_params(), input_ids, attention_mask))
+        return (output, None) if cache else output
+
+    def loss_and_backward(self, input_ids, target, attention_mask=None):
+        mx = _mlx()
+        target = mx.array(target, dtype=mx.int32)
+        params = self._mlx_params()
+
+        def loss_fn(current):
+            logits = self._forward_mlx(current, input_ids, attention_mask)
+            flat_logits = logits.reshape((-1, logits.shape[-1])).astype(mx.float32)
+            flat_target = target.reshape((-1,))
+            log_probs = flat_logits - mx.logsumexp(flat_logits, axis=-1, keepdims=True)
+            losses = -mx.take_along_axis(log_probs, flat_target[:, None], axis=1).squeeze(1)
+            if self.task == "text-generation":
+                valid = flat_target != self.pad_id
+                return mx.sum(mx.where(valid, losses, 0.0)) / mx.maximum(mx.sum(valid), 1)
+            return mx.mean(losses)
+
+        loss_scale = 128.0 if self.precision == "fp16" else 1.0
+        scaled_loss_fn = lambda current: loss_fn(current) * loss_scale
+        loss, gradients = mx.value_and_grad(scaled_loss_fn)(params)
+        for name, parameter in self.named_parameters():
+            parameter.grad[...] = np.asarray(gradients[name], dtype=np.float32) / loss_scale
+        return float(np.asarray(loss)) / loss_scale
+
+    def generate(self, input_ids, max_new_tokens, temperature=.8, top_k=40, eos_id: Optional[int] = None):
+        ids = np.asarray(input_ids, dtype=np.int64).copy()
+        for _ in range(max_new_tokens):
+            logits = self.forward(ids[:, -self.max_seq_len:])[:, -1, :] / max(temperature, 1e-5)
+            if top_k is not None:
+                k = min(top_k, logits.shape[-1])
+                excluded = np.argpartition(logits, -k, axis=1)[:, :-k]
+                logits[np.arange(len(ids))[:, None], excluded] = -np.inf
+            probs = np.exp(logits - logits.max(1, keepdims=True))
+            probs /= probs.sum(1, keepdims=True)
+            next_ids = np.array([np.random.choice(logits.shape[1], p=p) for p in probs])[:, None]
+            ids = np.concatenate([ids, next_ids], axis=1)
+            if eos_id is not None and np.all(next_ids == eos_id):
+                break
+        return ids