Metadata-Version: 2.3
Name: deciax
Version: 0.1.1
Summary: High-performance decision models on TPUs, powered by JAX and Pallas.
Keywords: decision-models,jax,pallas,tpu
Author: Junyi Hou
Author-email: Junyi Hou <hji200914@gmail.com>
Requires-Dist: huggingface-hub>=1.0
Requires-Dist: jax>=0.11.2
Requires-Dist: numpy>=2.3
Requires-Dist: safetensors>=0.6
Requires-Dist: tokenizers>=0.22
Requires-Dist: jax[tpu]>=0.11.2 ; extra == 'tpu'
Requires-Python: >=3.14
Project-URL: Repository, https://github.com/Junyi-99/deciax
Project-URL: Issues, https://github.com/Junyi-99/deciax/issues
Provides-Extra: tpu
Description-Content-Type: text/markdown

# deciax

High-performance decision models on TPUs, powered by JAX and Pallas.

## Development

Install the development tools and Git hooks:

```sh
uv sync
uv run pre-commit install
```

Run all checks manually with `uv run pre-commit run --all-files`.

Run the tests and type checks:

```sh
uv run pytest            # TPU-only tests are skipped off-TPU
uv run mypy
```

Parity fixtures come from the PyTorch reference, Laya 0.4.4 at commit `682ac67`. Install it with
`uv sync --group reference` and run `scripts/export_reference.py`.

## Layout

| Path | Contents |
|---|---|
| `src/deciax/protocol/` | Model-agnostic question/answer types and the `/v1/systemone` schema. Pure Python. |
| `src/deciax/models/` | Checkpoint registry, the `DecisionModel` interface, the ModernBERT encoder. |
| `src/deciax/models/laya/` | Laya: config, tokenizer, sequence construction, heads, calibration, weights. |
| `src/deciax/nn.py` | Backend-agnostic primitives (linear, layer norm, RoPE). |
| `src/deciax/ops/` | The only place the backend forks (`jax` reference or `pallas`). |
| `src/deciax/kernels/` | Attention: pure-jnp reference and the Pallas TPU kernel. |
| `src/deciax/runtime/` | Shape buckets and the jit engine. |
| `src/deciax/serve/` | `python -m deciax.serve`: stdlib HTTP server. |

## Usage

```python
import deciax

decider = deciax.load("laya")  # or "laya-multilingual"; bf16 on TPU, fp32 elsewhere
decider.decide(
    {"user": "I was charged twice this month, please fix it."},
    {
        "intent": {
            "type": "choice",
            "instructions": "What does the user want?",
            "criteria": {"billing": "payments, refunds", "other": "everything else"},
        },
        "urgency": {
            "type": "score",
            "instructions": "How urgent is this?",
            "criteria": ["not urgent", "soon", "blocking"],
        },
        "angry": {"type": "noul", "instructions": "Is the user angry?"},
    },
)
```

The response has the same shape as Laya's `system_one` response. Install with `uv sync --extra tpu` on a TPU host.

### HTTP server

```sh
uv run python -m deciax.serve --model laya --model laya-multilingual --default-model laya-multilingual --port 8000
curl -s localhost:8000/v1/systemone -d '{"state": "I was charged twice", "questions": {"angry": {"type": "noul", "instructions": "Is the user angry?"}}}'
```

The server speaks the same `/v1/systemone` wire protocol as `laya-serve`:
- **Routes**: `POST /v1/systemone`, `POST /v1/systemone/batch` and `GET /health`.
- **Requests and responses**: the same request limits and status codes, `x_jev_confidence` on choice and
  score answers, `--jev-strict` and Bearer auth (`DECIAX_API_KEY`).
- **Request controls**: `model`, `max_len` and `head_max_len`.

There are two differences from `laya-serve`:
- No language routing. A request with no `model`, the bundle id or an unknown id such as `jev-1` gets the
  default model.
- The controls deciax does not implement return 422 instead of being ignored: `task`, `lang`,
  `lang_guess`, `min_confidence`, `batch_size` and `sort_by_length`.

Every flag also reads an environment variable: `DECIAX_HOST`, `DECIAX_PORT`, `DECIAX_MODELS`,
`DECIAX_DEFAULT_MODEL`, `DECIAX_MAX_CONCURRENT`, `DECIAX_JEV_STRICT`, `DECIAX_DTYPE`, `DECIAX_IMPL` and
`DECIAX_LOG_LEVEL`.

## Status

The Laya inference path works end to end. The runtime needs only JAX and `tokenizers`; it does not need
torch, transformers or laya. Parity against PyTorch Laya 0.4.4 (`tests/laya/`):
- **Inputs**: tokenizer ids and `input_ids` / `marker_pos` match token for token.
- **CPU fp32**: option and act logits are within 1e-3 (act logits within 1e-5 relative). Responses match
  up to the last rounded digit.
- **TPU bf16**: probabilities are within 0.02 and the argmax matches, with both the `jax` and `pallas`
  attention.

`impl="auto"` uses XLA attention, which is faster than the Pallas flash kernel at these shapes.

Not yet supported:
- the parallel option layout
- per-language temperatures and binning maps
- `predict_long` windowing
- tournament selection
- language routing between checkpoints
- benchmarks

Laya is Apache-2.0 (https://github.com/NandhaKishorM/laya). Code ported from it is attributed in each
module docstring.
