class SinusoidalPositionalEmbedding(HybridBlock):
def __init__(self, units: int, dtype: Union[str, type] = 'float32'):
"""Use a geometric sequence of timescales.
Parameters
----------
units
The number of units for positional embedding
dtype
The dtype of the inner positional embeddings
"""
super().__init__()
def _init_sinusoidal_base(units):
half_units = units // 2
val = np.log(10000) / (half_units - 1)
val = np.exp(np.arange(half_units, dtype=np.float32) * -val)
return val
self._units = units
self._dtype = dtype
sinusoidal_base = _init_sinusoidal_base(units)
print(units, sinusoidal_base)
self.base_mult = Constant(sinusoidal_base)
def forward(self, positions):
"""
Parameters
----------
positions : NDArray
Shape (..., )
Returns
-------
ret :
Shape (..., units)
"""
emb = np.expand_dims(positions.astype(self._dtype),
axis=-1) * self.base_mult.data()
sin_emb = np.sin(emb)
cos_emb = np.cos(emb)
if self._units % 2 == 0:
return np.concatenate([sin_emb, cos_emb], axis=-1)
else:
return np.concatenate([
sin_emb, cos_emb,
np.expand_dims(np.zeros_like(positions).astype(self._dtype),
axis=-1)
],
axis=-1)
def __repr__(self):
s = '{name}(units={units}, dtype={dtype})'
return s.format(name=self.__class__.__name__,
units=self._units,
dtype=self._dtype)
def __init__(self, units: int, dtype: Union[str, type] = 'float32'):
"""Use a geometric sequence of timescales.
Parameters
----------
units
The number of units for positional embedding
dtype
The dtype of the inner positional embeddings
"""
super().__init__()
def _init_sinusoidal_base(units):
half_units = units // 2
val = np.log(10000) / (half_units - 1)
val = np.exp(np.arange(half_units, dtype=np.float32) * -val)
return val
self._units = units
self._dtype = dtype
sinusoidal_base = _init_sinusoidal_base(units)
self.base_mult = Constant(sinusoidal_base)
def __init__(self):
super().__init__()
self.weight = Constant(np.ones((10, 10)))