Handle invalid dimensions in SinusoidalPositionalEncoding

[pchintar]

Description

nn.SinusoidalPositionalEncoding currently assumes that dims is positive and even, but this requirement is not enforced.

The implementation constructs sinusoidal embeddings from sine/cosine frequency pairs using:

one_zero = 1 - mx.arange(0, dims // 2) / (dims // 2 - 1)

This works for the existing tested case (dims=16), but some edge cases are not handled correctly.

For example, if a user constructs:

nn.SinusoidalPositionalEncoding(3)

the implementation produces only 2 embedding dimensions instead of the requested 3. Similarly, dims=2 results in a division by zero while computing the values used to generate the sinusoidal frequencies, and non-positive values are not rejected explicitly.

This change validates that dims is a positive even integer and handles the dims=2 case separately to avoid the division by zero.

Tests

Added coverage for:

• dims=2
• Invalid dimensions: 0, 1, and 3

cc @zcbenz

[zcbenz]

Thanks!

[pchintar]

Thanks!

thank you @zcbenz