Fix LoRA embedding parameter initialization

[Aditya Biranje (AdityaBiranje)]

Description

Fixes the LoRA parameter initialization in Embedding to match the initialization used by Linear.

Previously, Embedding initialized lora_A to zeros and lora_B from a normal distribution. This was inconsistent with the expected LoRA initialization and with the existing Linear implementation.

The change initializes:

• lora_A with Kaiming uniform initialization
• lora_B with zeros

A regression test was added to verify the initialization behavior for both Embedding and Linear.

Testing

python -m pytest -q tests/test_layers.py
2 passed

The test was run locally on Apple Silicon using PyTorch 2.14.0.

Closes #194

[Aditya Biranje (AdityaBiranje)]

@microsoft-github-policy-service agree