tensorless 0.9.1__tar.gz → 0.9.2__tar.gz

{tensorless-0.9.1/tensorless.egg-info → tensorless-0.9.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tensorless
-Version: 0.9.1
+Version: 0.9.2
 Summary: ML with maximum automation and minimum setup.
 Author: Tensorless Contributors
 License: MIT
@@ -45,8 +45,11 @@ 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.
+CUDA, TPU, and MPS backends accelerate both transformer text tasks and
+tabular MLP tasks. Whatever device is auto-detected (or passed via
+`device=...`) is what training and inference actually run on; only
+unsupported platforms (no CUDA/TPU/MPS available) fall back to the native
+CPU engine.
 
 ## Train on your data
 

{tensorless-0.9.1 → tensorless-0.9.2}/README.md

@@ -19,8 +19,11 @@ 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.
+CUDA, TPU, and MPS backends accelerate both transformer text tasks and
+tabular MLP tasks. Whatever device is auto-detected (or passed via
+`device=...`) is what training and inference actually run on; only
+unsupported platforms (no CUDA/TPU/MPS available) fall back to the native
+CPU engine.
 
 ## Train on your data
 

tensorless-0.9.2/docs/installation.md

@@ -0,0 +1,49 @@
+# Installation
+
+## Requirements
+
+- Python 3.9 or later
+- NumPy 1.24 or later (installed automatically as a dependency)
+- Tensorless always runs its native vectorized engine, on CPU by default.
+  Installing the optional accelerator extras below lets it detect and use a
+  GPU/TPU automatically, for both transformer and tabular MLP models.
+
+## GPU / TPU / Apple Silicon support
+
+By default `pip install tensorless` only pulls in NumPy, so even on a
+machine with a GPU, device auto-detection will correctly report that no
+accelerator backend is *installed* and fall back to CPU. To actually train
+and run on your hardware, install the matching extra:
+
+```bash
+pip install 'tensorless[cuda]'   # NVIDIA GPUs, via JAX
+pip install 'tensorless[tpu]'    # Google TPUs, via JAX
+pip install 'tensorless[mps]'    # Apple Silicon, via MLX
+```
+
+Once installed, `tl.train(...)` auto-detects the best available device
+(`tpu` > `cuda` > `mps` > `cpu`) and trains on it without any extra
+configuration; you can also force a specific device with
+`tl.train(..., device="cuda")`.
+
+## Verify your installation
+
+```bash
+python -c "import tensorless as tl; print(tl.__version__)"
+tensorless --help
+```
+
+You should see a version string printed and the CLI's help text.
+
+## Runtime support
+
+Training and inference resolve the device automatically: the best available
+accelerator (`tpu` > `cuda` > `mps`) is used if its extra is installed and
+the hardware is detected, otherwise Tensorless falls back to the NumPy CPU
+engine. You can always force a specific device with `device="cpu"`,
+`device="cuda"`, etc. in `tl.train(...)`.
+
+## Troubleshooting installation
+
+See [troubleshooting.md](troubleshooting.md#installation-issues) for
+common installation problems.

{tensorless-0.9.1 → tensorless-0.9.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tensorless"
-version = "0.9.1"
+version = "0.9.2"
 description = "ML with maximum automation and minimum setup."
 readme = "README.md"
 requires-python = ">=3.9"

{tensorless-0.9.1 → tensorless-0.9.2}/tensorless/backends/jax_backend.py

@@ -12,6 +12,7 @@ import os
 import numpy as np
 
 from ..models.transformer import TinyTransformer
+from ..models.mlp import TabularMLP
 
 
 def _jax():
@@ -163,3 +164,75 @@ class JaxTinyTransformer(TinyTransformer):
             if eos_id is not None and np.all(next_ids == eos_id):
                 break
         return ids
+
+
+class JaxTabularMLP(TabularMLP):
+    """The native tabular-MLP parameter layout with JAX math and gradients.
+
+    Used automatically instead of the plain NumPy `TabularMLP` whenever the
+    resolved device is `cuda`/`tpu`, so tabular training actually runs on
+    the detected accelerator instead of silently staying on CPU.
+    """
+
+    def __init__(self, *args, **kwargs):
+        self.precision = kwargs.pop("precision", "fp32")
+        super().__init__(*args, **kwargs)
+
+    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, numeric, categorical, dropout_masks=None):
+        jax, jnp = _jax()
+        parts = [jnp.asarray(numeric, dtype=jnp.float32)] if self.n_numeric else []
+        categorical = jnp.asarray(categorical, dtype=jnp.int32)
+        for i in range(len(self.embeddings)):
+            parts.append(params[f"embeddings.{i}"][categorical[:, i]])
+        x = jnp.concatenate(parts, axis=1) if parts else jnp.asarray(numeric, dtype=jnp.float32)
+        for i in range(len(self.weights)):
+            preactivation = x @ params[f"weights.{i}"] + params[f"biases.{i}"]
+            x = jax.nn.gelu(preactivation, approximate=True)
+            if dropout_masks is not None and dropout_masks[i] is not None:
+                x = x * dropout_masks[i] / (1.0 - self.dropout)
+        out = x @ params["head_weight"] + params["head_bias"]
+        if self.task == "regression":
+            out = out[:, 0]
+        return out
+
+    def _make_dropout_masks(self, batch_size):
+        if not self.training or self.dropout <= 0:
+            return None
+        _, jnp = _jax()
+        return [
+            jnp.asarray((np.random.random((batch_size, weight.data.shape[1])) >= self.dropout).astype(np.float32))
+            for weight in self.weights
+        ]
+
+    def forward(self, numeric, categorical, cache=False):
+        params = self._jax_params()
+        batch_size = numeric.shape[0] if self.n_numeric else categorical.shape[0]
+        dropout_masks = self._make_dropout_masks(batch_size)
+        out = np.asarray(self._forward_jax(params, numeric, categorical, dropout_masks))
+        return (out, None) if cache else out
+
+    def loss_and_backward(self, numeric, categorical, target):
+        jax, jnp = _jax()
+        params = self._jax_params()
+        batch_size = numeric.shape[0] if self.n_numeric else categorical.shape[0]
+        dropout_masks = self._make_dropout_masks(batch_size)
+        is_regression = self.task == "regression"
+        target_arr = jnp.asarray(target, dtype=jnp.float32 if is_regression else jnp.int32)
+
+        def loss_fn(current):
+            logits = self._forward_jax(current, numeric, categorical, dropout_masks)
+            if is_regression:
+                error = logits - target_arr
+                return jnp.mean(error * error)
+            log_probs = jax.nn.log_softmax(logits, axis=-1)
+            return -jnp.mean(jnp.take_along_axis(log_probs, target_arr[:, None], axis=1))
+
+        loss, gradients = jax.value_and_grad(loss_fn)(params)
+        for name, parameter in self.named_parameters():
+            parameter.grad[...] = np.asarray(gradients[name], dtype=np.float32)
+        return float(loss)

{tensorless-0.9.1 → tensorless-0.9.2}/tensorless/backends/mlx_backend.py

@@ -11,6 +11,7 @@ from typing import Optional
 import numpy as np
 
 from ..models.transformer import TinyTransformer
+from ..models.mlp import TabularMLP
 
 
 def _mlx():
@@ -131,3 +132,81 @@ class MlxTinyTransformer(TinyTransformer):
             if eos_id is not None and np.all(next_ids == eos_id):
                 break
         return ids
+
+
+def _mlx_gelu(x, mx):
+    scale = float(np.sqrt(2.0 / np.pi))
+    scaled = scale * (x + 0.044715 * x ** 3)
+    return 0.5 * x * (1.0 + mx.tanh(scaled))
+
+
+class MlxTabularMLP(TabularMLP):
+    """The native tabular-MLP parameter layout with MLX math and gradients.
+
+    Used automatically instead of the plain NumPy `TabularMLP` whenever the
+    resolved device is `mps`, so tabular training actually runs on Apple
+    Silicon's GPU instead of silently staying on CPU.
+    """
+
+    def __init__(self, *args, **kwargs):
+        self.precision = kwargs.pop("precision", "fp32")
+        super().__init__(*args, **kwargs)
+
+    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, numeric, categorical, dropout_masks=None):
+        mx = _mlx()
+        parts = [mx.array(numeric, dtype=mx.float32)] if self.n_numeric else []
+        categorical = mx.array(categorical, dtype=mx.int32)
+        for i in range(len(self.embeddings)):
+            parts.append(params[f"embeddings.{i}"][categorical[:, i]])
+        x = mx.concatenate(parts, axis=1) if parts else mx.array(numeric, dtype=mx.float32)
+        for i in range(len(self.weights)):
+            preactivation = x @ params[f"weights.{i}"] + params[f"biases.{i}"]
+            x = _mlx_gelu(preactivation, mx)
+            if dropout_masks is not None and dropout_masks[i] is not None:
+                x = x * dropout_masks[i] / (1.0 - self.dropout)
+        out = x @ params["head_weight"] + params["head_bias"]
+        if self.task == "regression":
+            out = out[:, 0]
+        return out
+
+    def _make_dropout_masks(self, batch_size):
+        if not self.training or self.dropout <= 0:
+            return None
+        mx = _mlx()
+        return [
+            mx.array((np.random.random((batch_size, weight.data.shape[1])) >= self.dropout).astype(np.float32))
+            for weight in self.weights
+        ]
+
+    def forward(self, numeric, categorical, cache=False):
+        params = self._mlx_params()
+        batch_size = numeric.shape[0] if self.n_numeric else categorical.shape[0]
+        dropout_masks = self._make_dropout_masks(batch_size)
+        out = np.asarray(self._forward_mlx(params, numeric, categorical, dropout_masks))
+        return (out, None) if cache else out
+
+    def loss_and_backward(self, numeric, categorical, target):
+        mx = _mlx()
+        params = self._mlx_params()
+        batch_size = numeric.shape[0] if self.n_numeric else categorical.shape[0]
+        dropout_masks = self._make_dropout_masks(batch_size)
+        is_regression = self.task == "regression"
+        target_arr = mx.array(target, dtype=mx.float32 if is_regression else mx.int32)
+
+        def loss_fn(current):
+            logits = self._forward_mlx(current, numeric, categorical, dropout_masks)
+            if is_regression:
+                error = logits - target_arr
+                return mx.mean(error * error)
+            log_probs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
+            return -mx.mean(mx.take_along_axis(log_probs, target_arr[:, None], axis=1))
+
+        loss, gradients = mx.value_and_grad(loss_fn)(params)
+        for name, parameter in self.named_parameters():
+            parameter.grad[...] = np.asarray(gradients[name], dtype=np.float32)
+        return float(np.asarray(loss))

{tensorless-0.9.1 → tensorless-0.9.2}/tensorless/models/registry.py

@@ -53,7 +53,14 @@ def build_model(task: str, model_type: str, cfg: Dict[str, Any], meta: Dict[str,
             model_kwargs["gradient_checkpointing"] = cfg.get("gradient_checkpointing", False)
         return model_class(**model_kwargs)
     elif model_type == "mlp":
-        return TabularMLP(
+        model_class = TabularMLP
+        if backend == "jax":
+            from ..backends.jax_backend import JaxTabularMLP
+            model_class = JaxTabularMLP
+        elif backend == "mlx":
+            from ..backends.mlx_backend import MlxTabularMLP
+            model_class = MlxTabularMLP
+        model_kwargs = dict(
             n_numeric=meta["n_numeric"],
             categorical_vocab_sizes=meta["categorical_vocab_sizes"],
             d_model=cfg["d_model"],
@@ -62,5 +69,8 @@ def build_model(task: str, model_type: str, cfg: Dict[str, Any], meta: Dict[str,
             task=task,
             n_classes=meta.get("n_classes", 0),
         )
+        if backend in ("jax", "mlx"):
+            model_kwargs["precision"] = cfg.get("precision", "fp32")
+        return model_class(**model_kwargs)
     else:
         raise ModelError(f"Unknown model_type '{model_type}'.")

{tensorless-0.9.1 → tensorless-0.9.2}/tensorless/runtime.py

@@ -45,11 +45,10 @@ class LoadedModel:
         self.device = get_device(device_name)
 
         backend = "numpy"
-        if self.model_type == "transformer" and self.task in ("text-generation", "text-classification"):
-            if self.device in ("cuda", "tpu"):
-                backend = "jax"
-            elif self.device == "mps":
-                backend = "mlx"
+        if self.device in ("cuda", "tpu"):
+            backend = "jax"
+        elif self.device == "mps":
+            backend = "mlx"
         self.model = build_model(self.task, self.model_type, self.config, self.meta, backend=backend)
         self.model.load_state_dict(payload["model_state_dict"])
         self.model.eval()

{tensorless-0.9.1 → tensorless-0.9.2}/tensorless/training/trainer.py

@@ -66,12 +66,16 @@ def run_training(ds: Dataset, cfg: Dict[str, Any], checkpoint_mgr: CheckpointMan
     elif task == "text-classification": prepared = dp.prepare_text_classification(ds, cfg, tokenizer=tokenizer, classes=resume_state["meta"]["classes"] if resume_state else None)
     elif task in ("classification", "regression"): prepared = dp.prepare_tabular(ds, cfg, task, preprocessor)
     else: raise ValueError(f"Unsupported task '{task}'")
+    # Pick a backend that actually runs on the resolved device. Both model
+    # types (transformer and mlp) have jax/mlx implementations, so whatever
+    # device was auto-detected (or requested) is honored for training --
+    # we don't silently fall back to the NumPy/CPU engine just because the
+    # model happens to be an MLP.
     backend = "numpy"
-    if model_type == "transformer" and task in ("text-generation", "text-classification"):
-        if device in ("cuda", "tpu"):
-            backend = "jax"
-        elif device == "mps":
-            backend = "mlx"
+    if device in ("cuda", "tpu"):
+        backend = "jax"
+    elif device == "mps":
+        backend = "mlx"
     model = build_model(task, model_type, cfg, prepared.meta, backend=backend)
     if pretrained_state is not None:
         try:

{tensorless-0.9.1 → tensorless-0.9.2/tensorless.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tensorless
-Version: 0.9.1
+Version: 0.9.2
 Summary: ML with maximum automation and minimum setup.
 Author: Tensorless Contributors
 License: MIT
@@ -45,8 +45,11 @@ 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.
+CUDA, TPU, and MPS backends accelerate both transformer text tasks and
+tabular MLP tasks. Whatever device is auto-detected (or passed via
+`device=...`) is what training and inference actually run on; only
+unsupported platforms (no CUDA/TPU/MPS available) fall back to the native
+CPU engine.
 
 ## Train on your data
 

tensorless-0.9.1/docs/installation.md

@@ -1,29 +0,0 @@
-# Installation
-
-## Requirements
-
-- Python 3.9 or later
-- NumPy 1.24 or later (installed automatically as a dependency)
-- Tensorless currently runs its native vectorized engine on CPU; accelerator
-  backends are planned without changing the public API.
-
-## Verify your installation
-
-```bash
-python -c "import tensorless as tl; print(tl.__version__)"
-tensorless --help
-```
-
-You should see a version string printed and the CLI's help text.
-
-## Runtime support
-
-The native engine uses NumPy vectorized CPU kernels and automatically resolves
-the device to CPU. You can still specify `device="cpu"` explicitly in
-`tl.train(...)`; the device option remains forward-compatible with future
-accelerator backends.
-
-## Troubleshooting installation
-
-See [troubleshooting.md](troubleshooting.md#installation-issues) for
-common installation problems.