[Flax] time embedding (#1081)

Commit

2 years ago

[Flax] time embedding (#1081) * initial get_sinusoidal_embeddings * added asserts * better var name * fix docs