fix(embeddings): slice Matryoshka dimensions on final axis

[linjiaye929]

Summary

Matryoshka slicing currently treats the first axis of each output as its embedding width. For ColBERT-style outputs shaped (tokens, embedding_width), this truncates tokens, leaves the embedding width unchanged, and validates requested dimensions against the first input's token count.

Validate against shape[-1] and slice with [..., :matryoshka_dim]. Dense embeddings keep their existing behavior. Late-interaction embeddings retain every token while selecting the requested leading coordinates.

Related to #559: this addresses the reported dimension-axis bug only. Query/document-specific processing, projection, and reranking behavior remain outside this change. This PR does not close that issue.

Validation

Linux, Python 3.12.14, NumPy 1.26.4, CPU-only PyTorch 2.8.0, sentence-transformers 3.3.1, and pytest 8.4.2.

From libs/infinity_emb:

DO_NOT_TRACK=1 HF_HUB_OFFLINE=1 python -m pytest tests/unit_test/inference/test_matryoshka.py -q

• 21 passed after the change.
• The same tests on unchanged 040269b give 10 failed, 11 passed.
• Coverage includes dense embeddings, variable-length token embeddings in both input orders, dimensions below and above token counts, full-width slicing, invalid dimensions, unchanged None/0 behavior, dtype preservation, and non-mutation of inputs.
• The ColBERT fixture invokes Infinity's actual SentenceTransformerPatched.encode_post with deterministic PyTorch tensors and attention masks. A SimpleNamespace supplies only postprocessor configuration; it is not a loaded model instance. No model weights are loaded, and these tests make no pretrained-model or reduced-dimension retrieval-quality claim.

Adjacent checks:

DO_NOT_TRACK=1 HF_HUB_OFFLINE=1 python -m pytest \
  tests/unit_test/inference/test_matryoshka.py \
  tests/unit_test/test_engine.py::test_async_api_debug \
  tests/unit_test/test_engine.py::test_args_between_array_and_engine_same \
  tests/unit_test/test_sync_engine.py::test_sync_engine \
  tests/unit_test/test_sync_engine.py::test_sync_engine_double \
  tests/unit_test/inference/test_select_model.py \
  -k 'not torch and not optimum and not ctranslate2' -q
python -m pytest tests/unit_test/test_args.py -q

• 30 passed, 3 deselected, then 3 passed.
• ruff check . passes. Changed-file formatting, codespell, and git diff --check pass.
• mypy infinity_emb reports the same 22 missing-optional-dependency errors on this branch and unchanged main, with identical output and no batch_handler.py errors.
• Full make precommit, pretrained-model, accelerator, and HTTP-server suites were not run in this minimal CPU environment.

Checklist

• Read the contribution guidelines.
• Added regression tests.
• Documented embedding-dimension slicing.

AI assistance was used to prepare this contribution.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-12T10:34:28.103013Z	bba670d	Draft marked ready

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[greptile-apps]

Confidence Score: 5/5

[Medium impact] This PR appears safe to merge; it fixes the dimension axis without changing dense embedding behavior.

Summary

Fixes embedding-width checks and slicing to use the final axis.

• Matryoshka slicing shortens embedding width without dropping tokens.

linjiaye929 explicitly leaves query/document-specific processing, projection, and reranking outside this change. This PR does not claim to close #559.

No actionable issues were found. Tests were inspected, not run during this review.

_{Reviews (1) · Last reviewed commit: "fix(embeddings): slice Matryoshka dimens..." · Reviewed by Greptile}