def _apply_1d(self, input_slice: jax.Array, position: jax.Array) -> jax.Array:
    """Apply RoPE to a one-dimensional JAX array."""
    assert input_slice.ndim == 1
    assert position.ndim == 0
    # Infer `head_dim` from the input shape
    [head_dim] = input_slice.shape
    fraction = 2 * jnp.arange(0, head_dim // 2) / head_dim
    timescale = self.max_wavelength**fraction
    # Since we're assuming `timescale` is a vector and `position` is a scalar,
    # we don't need any axis alignment.
    sinusoid_inp = position / timescale
    sin = jnp.sin(sinusoid_inp)
    cos = jnp.cos(sinusoid_inp)
    first_half, second_half = jnp.split(input_slice, 2)
    first_part = first_half * cos - second_half * sin
    second_part = second_half * cos + first_half * sin
    return jnp.concatenate([first_part, second_part])
Apply RoPE to a one-dimensional JAX array.
penzai/deprecated/v1/nn/embeddings.py