def _sincos(self, start, length, d_feature):
    """Create the sin-cos tensor of shape [1, length, d_feature]."""
    position = jnp.arange(0, length)[:, None] + start
    div_term = jnp.exp(
        jnp.arange(0, d_feature, 2) * -(jnp.log(10000.0) / d_feature))
    sin = jnp.sin(position * div_term)
    cos = jnp.cos(position * div_term)
    pe = jnp.concatenate([sin, cos], axis=1)
    return pe[None, :, :]  # [1, length, d_feature]
Create the sin-cos tensor of shape [1, length, d_feature].
trax/layers/research/position_encodings.py