[Afloat16]

Problem

StableEmbedding passes device and dtype positionally to torch.nn.Embedding. The parent signature now places _freeze before these factory arguments, so an explicit device silently freezes the embedding weight, while an explicit dtype is interpreted as a device and raises a TypeError. Meta-device initialization also allocates the embedding on CPU instead of meta.

Change

Forward the two existing factory arguments by keyword. No public arguments change. The normalization parameters retain their existing full-precision initialization.

Reproducer

import torch
from bitsandbytes.nn import StableEmbedding

embedding = StableEmbedding(32, 16, device="cpu")
assert embedding.weight.requires_grad
embedding = StableEmbedding(32, 16, device="cpu", dtype=torch.bfloat16)
embedding(torch.tensor([1, 2, 3])).float().square().mean().backward()

Unchanged main fails the first assertion; constructing the explicitly typed layer separately raises TypeError: empty(): argument 'device' must be torch.device, not torch.dtype.

Validation

Base: a616886cacf86701514f0cb1d6be760f967481c0; tested patch: 42530450deaaef630f95a89bbeaac07a9118177c.

• The same 23 new cases: 22 failed, 1 passed before; 23 passed after.
• New tests plus the complete existing tests/test_modules.py on CPU: 110 passed, 2 skipped.
• Fifteen device/dtype combinations compare an Embedding + LayerNorm + classifier reference through three momentum-SGD updates, including outputs, all gradients, parameters, optimizer buffers and the padding row.
• Additional cases cover meta allocation, explicit pretrained weights, max-norm clipping, frequency-scaled gradients, sparse gradients and negative padding indices.
• Built the CPU native library from current source; Python 3.12.14 / PyTorch 2.15.0+cpu. Repository-pinned Ruff 0.14.3 checks/formatting and git diff --check pass.

BNB_TEST_DEVICE=cpu OMP_NUM_THREADS=2 python -m pytest tests/test_stable_embedding_constructor.py tests/test_modules.py -q

GPU/distributed execution and the complete repository suite were not run.

Checked all-state constructor/device/dtype/freeze issues and PRs, the module history, the latest 100 PRs and the user fork. #87 introduced the positional arguments; it does not address the later parent-signature change. #2132 fixes the separate forward normalization dtype problem. This patch is independently based on main and does not include or require #2132.