def _get_embeddings(self, t):
    """Get embeddings float[..., num_features].

    Args:
      t: int[...] position (i.e. jnp.arange(..., jnp.int32))

    Returns:
      embeddings: float[..., num_features]
    """
    inter_bin_idx, intra_bin_idx = divmod(t, self._time_bin_length)
    bin_parity = inter_bin_idx % 2
    bin_fraction = intra_bin_idx / self._time_bin_length
    embeddings = jnp.stack([
        1 / (1 + inter_bin_idx),
        bin_fraction,
        bin_parity.astype(jnp.float32),
    ], -1)

    assert embeddings.shape == t.shape + (self.num_features,), embeddings.shape
    return embeddings
Get embeddings float[..., num_features].

    Args:
      t: int[...] position (i.e. jnp.arange(..., jnp.int32))

    Returns:
      embeddings: float[..., num_features]
    
trax/layers/research/position_encodings.py