[flax_llama] Why is the return value of the `create_sinusoidal_positions` truncated by `num_pos`?

mantle2048 commented 5 months ago

System Info

transformers version: 4.38.2
Platform: Linux-3.10.0-1160.el7.x86_64-x86_64-with-glibc2.35
Python version: 3.12.2
Huggingface_hub version: 0.21.3
Safetensors version: 0.4.2
Accelerate version: not installed
Accelerate config: not found
PyTorch version (GPU?): not installed (NA)
Tensorflow version (GPU?): not installed (NA)
Flax version (CPU?/GPU?/TPU?): 0.8.1 (gpu)
Jax version: 0.4.23
JaxLib version: 0.4.23
Using GPU in script?: Yes
Using distributed or parallel set-up in script?: Yes

Who can help?

No response

Information

[ ] The official example scripts
[ ] My own modified scripts

Tasks

[ ] An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
[ ] My own task or dataset (give details below)

Reproduction

I noticed that in the implementation of flax-llama, the return value of create_sinusoidal_positions was truncated by num_pos.

def create_sinusoidal_positions(num_pos, dim):
    inv_freq = 1.0 / (10000 ** (np.arange(0, dim, 2) / dim))
    freqs = np.einsum("i , j -> i j", np.arange(num_pos), inv_freq).astype("float32")

    emb = np.concatenate((freqs, freqs), axis=-1)
    out = np.concatenate((np.sin(emb)[:, None, :], np.cos(emb)[:, None, :]), axis=-1)
    return jnp.array(out[:, :, :num_pos]) # Why?

Why is this? Is it to prevent the sequence length from being too short?

Expected behavior

See above.

amyeroberts commented 5 months ago

cc @sanchit-gandhi

amyeroberts commented 4 months ago

Gentle ping @sanchit-gandhi

amyeroberts commented 3 months ago

Another ping @sanchit-gandhi

amyeroberts commented 1 month ago

Another another ping @sanchit-gandhi. Could you nominate someone to take over this whilst you're away?

huggingface / transformers