# Test Phoenix

> uv run pytest targets, attention-mask and retrieval assertions, and the success criteria encoded in the Phoenix test modules.

- Repository: xai-org/x-algorithm
- GitHub: https://github.com/xai-org/x-algorithm
- Human docs: https://grok-wiki.com/public/docs/xai-org-x-algorithm-23c09c39074c
- Complete Markdown: https://grok-wiki.com/public/docs/xai-org-x-algorithm-23c09c39074c/llms-full.txt

## Source Files

- `phoenix/test_recsys_model.py`
- `phoenix/test_recsys_retrieval_model.py`
- `phoenix/pyproject.toml`
- `phoenix/README.md`
- `phoenix/grok.py`
- `phoenix/recsys_model.py`

---

---
title: "Test Phoenix"
description: "uv run pytest targets, attention-mask and retrieval assertions, and the success criteria encoded in the Phoenix test modules."
---

Phoenix unit tests live in `phoenix/test_recsys_model.py` and `phoenix/test_recsys_retrieval_model.py`. From `phoenix/`, the documented target is `uv run pytest` on those two files. The suite asserts `make_recsys_attn_mask` candidate isolation, `right_anchored_rope_positions`, `compute_post_age_bucket`, `normalize_continuous_value`, `CandidateTower` L2 outputs, and synthetic retrieval `RetrievalOutput` shapes. It initializes Haiku params on random `create_example_batch` / `create_example_corpus` tensors and does **not** load `oss-phoenix-artifacts`.

<Warning>
These tests are not the inference-pipeline success signal. A ranked table from `run_pipeline.py` requires extracted checkpoints. A green pytest run only means the JAX helpers and a randomly initialized retrieval model still satisfy the contracts below.
</Warning>

## Prerequisites

<ParamField body="requires-python" type="string" required>
`>=3.11` from `phoenix/pyproject.toml`.
</ParamField>

<ParamField body="runtime deps" type="list">
`dm-haiku>=0.0.13`, `jax==0.8.1`, `numpy>=1.26.4`. `pyright` is also a project dependency; the test modules do not invoke it.
</ParamField>

<ParamField body="dev deps" type="list">
`pytest` is listed under `[dependency-groups] dev`, not under `[project] dependencies`.
</ParamField>

<ParamField body="uv environments" type="list">
`[tool.uv] environments` is restricted to `sys_platform == 'darwin'` and `sys_platform == 'linux'`.
</ParamField>

There is no `pytest.ini` and no `[tool.pytest]` table. There is no CI workflow in this checkout that runs these files. Modules are flat (`from grok import ...`, `from recsys_model import ...`), so the working directory must be `phoenix/` or that directory must be on `PYTHONPATH`.

:::files
phoenix/
  test_recsys_model.py              pytest-style classes: mask, RoPE, age, norms
  test_recsys_retrieval_model.py    unittest.TestCase: tower, model, runner
  grok.py                           make_recsys_attn_mask, right_anchored_rope_positions
  recsys_model.py                   compute_post_age_bucket, NormConfig, HashConfig
  recsys_retrieval_model.py         CandidateTower, PhoenixRetrievalModel, RetrievalOutput
  runners.py                        create_example_batch, create_example_corpus, RecsysRetrievalInferenceRunner
  pyproject.toml                    pytest in [dependency-groups] dev
:::

## Run the suite

<Steps>
<Step title="Install from phoenix/">
```bash
cd phoenix
uv sync
```

`uv sync` installs the project plus the default `dev` group, which provides `pytest`.
</Step>

<Step title="Run the two named files">
```bash
uv run pytest test_recsys_model.py test_recsys_retrieval_model.py
```

That is the command in `phoenix/README.md`. From the repo root, prefix the path: `uv run --directory phoenix pytest test_recsys_model.py test_recsys_retrieval_model.py` only works if imports still resolve; prefer `cd phoenix`.
</Step>

<Step title="Confirm every collected method passes">
Success is pytest exit code `0`. The two files encode **22** pytest methods in `test_recsys_model.py` and **12** `unittest.TestCase` methods in `test_recsys_retrieval_model.py`. Failures are assertion messages such as `Position i should NOT attend to future position j` or a shape mismatch on `user_representation`.
</Step>
</Steps>

<CodeGroup>
```bash title="Documented pytest target"
cd phoenix
uv run pytest test_recsys_model.py test_recsys_retrieval_model.py
```

```bash title="Verbose / one class"
cd phoenix
uv run pytest test_recsys_model.py::TestMakeRecsysAttnMask -v
uv run pytest test_recsys_retrieval_model.py::TestPhoenixRetrievalModel::test_retrieve_top_k -v
```

```bash title="File __main__ runners"
cd phoenix
uv run python test_recsys_model.py
uv run python test_recsys_retrieval_model.py
```
</CodeGroup>

`test_recsys_model.py` ends with `pytest.main([__file__, "-v"])`. `test_recsys_retrieval_model.py` ends with `unittest.main()`. pytest collects both styles when you pass the files.

<RequestExample>
```bash
cd phoenix && uv run pytest test_recsys_model.py test_recsys_retrieval_model.py -q
```
</RequestExample>

<ResponseExample>
```text
# Expected: process exit 0, every collected method in the two files passed.
# Not expected: artifact-path errors, sports_corpus.npz loads, or a ranked feed table.
```
</ResponseExample>

## What the suite covers

| File | Style | Classes | Surface under test |
|---|---|---|---|
| `test_recsys_model.py` | pytest classes | `TestMakeRecsysAttnMask`, `TestRightAnchoredRopePositions`, `TestComputePostAgeBucket`, `TestNormalizeContinuousValue` | Ranking-path helpers in `grok.py` and `recsys_model.py` |
| `test_recsys_retrieval_model.py` | `unittest.TestCase` | `TestCandidateTower`, `TestPhoenixRetrievalModel`, `TestRetrievalInferenceRunner` | `CandidateTower`, `PhoenixRetrievalModelConfig.make()`, `RecsysRetrievalInferenceRunner` |

The retrieval user tower concatenates user + history and calls the transformer with `candidate_start_offset=None`. Isolation-mask tests therefore apply to ranking (`make_recsys_attn_mask`), not to retrieval encoding.

## Attention-mask success criteria

`make_recsys_attn_mask(seq_len, candidate_start_offset, dtype=jnp.float32)` returns a `[1, 1, seq_len, seq_len]` array. `1` means attend; `0` means do not. Implementation starts from `jnp.tril`, zeros the candidate–candidate block, then restores the candidate diagonal.

<ParamField body="seq_len" type="int" required>
Total length: user prefix + history + candidates.
</ParamField>

<ParamField body="candidate_start_offset" type="int" required>
First candidate index. Positions `[0, candidate_start_offset)` are user + history.
</ParamField>

<ParamField body="dtype" type="jnp.dtype">
Default `jnp.float32`. `test_dtype_preserved` also constructs `jnp.float16` and asserts the dtype is kept.
</ParamField>

`TestMakeRecsysAttnMask` locks these rules:

| Method | Pass condition |
|---|---|
| `test_output_shape` | `mask.shape == (1, 1, seq_len, seq_len)` |
| `test_user_history_has_causal_attention` | For `i, j < candidate_start_offset`: `1` iff `j <= i` |
| `test_candidates_attend_to_user_history` | Every candidate row is `1` on all user/history keys |
| `test_candidates_attend_to_themselves` | Candidate diagonal is `1` |
| `test_candidates_do_not_attend_to_other_candidates` | Off-diagonal candidate–candidate cells are `0` |
| `test_full_mask_structure` | Exact matrix for `[user, h1, h2, c1, c2, c3]` |
| `test_single_candidate` | Causal prefix plus one self-attending candidate |
| `test_all_candidates` | Offset `1`: only position `0` is the prefix; each later token attends to user + self |

Canonical fixture (`seq_len=6`, `candidate_start_offset=3`):

```text
# keys →
#          u  h1 h2 c1 c2 c3
user    [  1, 0, 0, 0, 0, 0 ]
h1      [  1, 1, 0, 0, 0, 0 ]
h2      [  1, 1, 1, 0, 0, 0 ]
c1      [  1, 1, 1, 1, 0, 0 ]
c2      [  1, 1, 1, 0, 1, 0 ]
c3      [  1, 1, 1, 0, 0, 1 ]
```

<Note>
User + history attention in this function is **causal** (`jnp.tril`), not bidirectional. `test_user_history_has_causal_attention` fails if a history token attends to a later history token. That is the contract to keep when editing `make_recsys_attn_mask`.
</Note>

The all-candidates edge (`candidate_start_offset=1`) is the isolation rule with no history: each candidate attends to the user token and itself only.

## RoPE, age buckets, and continuous norms

### `right_anchored_rope_positions`

Signature: `(padding_mask, history_seq_len, num_user_prefix_tokens) -> [B, T]`.

| Method | Pass condition |
|---|---|
| `test_output_shape` | Shape `[B, T]` |
| `test_prefix_positions_preserved` | Prefix slots are `0 .. num_user_prefix_tokens-1` |
| `test_candidates_share_position` | Every candidate index equals `history_end = num_prefix + history_seq_len` |
| `test_padding_gets_zero` | `padding_mask == False` positions are `0.0` |

Newest history is right-anchored so the last valid history token keeps a fixed RoPE index. Tokens after `history_end` (candidates) share that `history_end` index.

### `compute_post_age_bucket`

`POST_AGE_MAX_MINUTES = 4800`. Default `granularity_mins=60`. Formula: `bucket = (age_minutes // granularity_mins) + 1`, clipped to the overflow bucket `4800 // granularity_mins + 1`, then forced to `0` when age is negative or either timestamp is `0`.

| Method | Input (minutes old, 60-min buckets) | Expected bucket |
|---|---|---|
| `test_basic_bucketing` | 30 | `1` (0–59) |
| `test_two_hour_post` | 120 | `3` (120–179) |
| `test_missing_timestamp_zero` | `impr_ts == 0` or `post_ts == 0` | `0` |
| `test_negative_age_maps_to_zero` | post created after impression | `0` |
| `test_overflow_bucket` | 5000 | `81` (`4800 // 60 + 1`) |
| `test_batch_processing` | `[30, 120, missing]` | `[1, 3, 0]`, shape `(1, 3)` |

### `normalize_continuous_value`

`NormConfig(norm_scale=30.0, use_log=False)` by default. Values clip to `[0, norm_scale]` first.

| Method | Config | Input | Expected |
|---|---|---|---|
| `test_linear_normalization` | `norm_scale=30`, `use_log=False` | `[0, 15, 30, 60]` | `[0.0, 0.5, 1.0, 1.0]` |
| `test_log_normalization` | `norm_scale=30`, `use_log=True` | `[0, 30]` | `[0.0, 1.0]` (`log1p(x) / log1p(scale)`) |
| `test_clamping` | `norm_scale=10`, linear | `[-5, 0, 5, 15]` | `[0.0, 0.0, 0.5, 1.0]` |

Linear path: `clip(x, 0, scale) / scale`. Log path: `log1p(clip(x, 0, scale)) / log1p(scale)`.

## Retrieval success criteria

Retrieval tests construct a **tiny** `PhoenixRetrievalModelConfig` (not the published mini checkpoint): `emb_size=64`, `history_seq_len=16`, `candidate_seq_len=8`, `batch_size=2`, `num_actions=19`, `corpus_size=100`, `top_k=10`, `HashConfig(num_user_hashes=2, num_item_hashes=2, num_author_hashes=2)`, `product_surface_vocab_size=16`, and a 1-layer `TransformerConfig` (`widening_factor=2`, `key_size=32`, `num_q_heads=2`, `num_kv_heads=2`, `attn_output_multiplier=0.125`). Forward graphs use `hk.without_apply_rng(hk.transform(...))` and `jax.random.PRNGKey(0)`.

### `CandidateTower`

Input shape in the tests: `[B, C, num_hashes, D]` with `B=4`, `C=8`, `num_hashes=4`, `D=64`.

| Method | `enable_linear_proj` | Pass condition |
|---|---|---|
| `test_candidate_tower_output_shape` | `True` | Output `[B, C, D]` |
| `test_candidate_tower_normalized` | `True` | Per-vector L2 norm `≈ 1` (`decimal=5`) |
| `test_candidate_tower_mean_pooling` | `False` | Same shape and unit L2 |
| `test_mean_pooling_has_no_params` | `False` | `sum(p.size for p in jax.tree.leaves(params)) == 0` |

Linear mode is a two-layer SiLU MLP (`candidate_tower_projection_1`, `candidate_tower_projection_2`) then L2. Mean-pool mode averages on the hash axis, then L2.

### `PhoenixRetrievalModel`

`__call__(batch, embeddings, corpus_embeddings, top_k)` returns `RetrievalOutput`.

<ResponseField name="user_representation" type="jax.Array">
`[B, D]`. `test_user_representation_normalized` requires unit L2 per row.
</ResponseField>

<ResponseField name="top_k_indices" type="jax.Array">
`[B, K]`. Every index in `[0, corpus_size)`.
</ResponseField>

<ResponseField name="top_k_scores" type="jax.Array">
`[B, K]`. Per-row scores are nonincreasing (`scores[i] >= scores[i+1]`).
</ResponseField>

| Method | Pass condition |
|---|---|
| `test_model_forward` | The three `RetrievalOutput` shapes above |
| `test_candidate_representation_normalized` | `build_candidate_representation` → `[B, C, D]` unit L2 |
| `test_retrieve_top_k` | In-range indices and descending scores |
| `test_mean_pooling_model_forward` | Same user / top-k shapes with `enable_linear_proj=False` |

Scoring is `user_representation @ corpus_embeddings.T` then `jax.lax.top_k`. Corpus rows from `create_example_corpus` are already L2-normalized.

### `RecsysRetrievalInferenceRunner`

`initialize()` builds dummy batch/embeddings from the config, a dummy corpus of shape `(10, emb_size)`, and `dummy_top_k=5`, then stores `runner.params`. Tests pass `RetrievalModelRunner(..., bs_per_device=0.125)` (`BaseModelRunner` default is `2.0`).

| Method | Pass condition |
|---|---|
| `test_runner_initialization` | `runner.params is not None` after `initialize()` |
| `test_runner_encode_user` | `encode_user` → `[B, D]` |
| `test_runner_retrieve` | After `set_corpus(corpus_embeddings, corpus_post_ids)`, `retrieve(..., top_k=10)` matches `RetrievalOutput` shapes |

## Synthetic fixtures

`create_example_batch` (seed `42`) and `create_example_corpus` (seed `123`) are the only data sources. Hash `0` is reserved for padding; generated hashes are in `[1, num_*_embeddings)`. History rows are randomly truncated to `[history_len // 2, history_len]`. Actions are Bernoulli with `p=0.3` (`> 0.7` kept). Corpus embeddings are Gaussian then L2-normalized; `corpus_post_ids` is `arange(corpus_size)`.

<AccordionGroup>
<Accordion title="Fixture sizes used by TestPhoenixRetrievalModel">

| Knob | Test value | Notes |
|---|---|---|
| `emb_size` | `64` | Not the published mini-model width |
| `history_seq_len` | `16` | Published mini config uses `127` |
| `candidate_seq_len` | `8` | Published mini config uses `64` |
| `num_actions` | `19` | Matches `BaseInferenceRunner._get_num_actions` fallback |
| `corpus_size` | `100` | Not `sports_corpus.npz` |
| `top_k` | `10` | Not `--top_k_retrieval` |
| Transformer | 1 layer, 2 heads | Not the exported ranker/retrieval depth |

</Accordion>
</AccordionGroup>

<Check>
Changing isolation, RoPE candidate sharing, overflow bucket `81`, unit-norm towers, or descending `top_k` without updating these assertions is a failed contract, not a style change.
</Check>

## What these tests do not cover

| Surface | Status |
|---|---|
| `run_pipeline.py` / `run_ranker.py` / `run_retrieval.py` | Not imported |
| `oss-phoenix-artifacts` (`model_params.npz`, `embedding_tables.npz`, `config.json`, `sports_corpus.npz`, `example_sequence.json`) | Not loaded |
| Ranking `PhoenixModelConfig` forward / multi-action logits | No test module |
| Hash-table lookup against the 1M-row published tables | Random embeddings only |
| Home Mixer, Thunder, Grox | Out of this directory |
| pyright | Dependency only |

A green unit suite plus a failed `run_pipeline.py` is consistent: tests never touch LFS artifacts.

## Troubleshooting

<AccordionGroup>
<Accordion title="ModuleNotFoundError: grok / recsys_model / recsys_retrieval_model">
Run from `phoenix/`. The tree is not an installed package; pytest adds the current directory to `sys.path`.
</Accordion>

<Accordion title="pytest: command not found inside uv">
`pytest` is a `dev` group extra. Re-run `uv sync` in `phoenix/`, then `uv run pytest ...` (not a bare `pytest` from a non-project venv).
</Accordion>

<Accordion title="JAX / Haiku install or platform errors">
`jax` is pinned to `0.8.1`. uv environments list only Darwin and Linux. Windows is not an advertised target.
</Accordion>

<Accordion title="Looking for artifact or viewer_id failures">
Those errors come from the inference CLI or Home Mixer, not these tests. See [Troubleshooting](/troubleshooting) and [Run the inference pipeline](/run-inference-pipeline).
</Accordion>
</AccordionGroup>

## Next

<CardGroup>
<Card title="Candidate isolation" href="/candidate-isolation">
Mask rules these tests freeze: candidates attend to user, history, and self only.
</Card>
<Card title="Phoenix model configuration" href="/phoenix-model-configuration">
`PhoenixRetrievalModelConfig`, `TransformerConfig`, `RecsysBatch`, and published `config.json` keys versus the tiny test config.
</Card>
<Card title="Run the inference pipeline" href="/run-inference-pipeline">
Checkpoint load, `sports_corpus.npz` retrieval, and the ranked table — a different success signal.
</Card>
<Card title="Installation" href="/installation">
Python 3.11+, uv, and Git LFS for artifacts (needed for inference, not for pytest).
</Card>
<Card title="Troubleshooting" href="/troubleshooting">
LFS, artifact paths, `viewer_id`, and other runtime failures outside the unit suite.
</Card>
</CardGroup>