import torch


def PositionEncoder(x):
	"""Positional Encoding layer with fixed encoding vector based on sin and cos,
	as in http://nlp.seas.harvard.edu/2018/04/03/attention.html#position-wise-feed-forward-networks
	Implemented as a function instead of a module because we may not know the shape of x
	(in particular, the text_len dimension) before hand.
	This function returns just the fixed PE vector instead of the sum x + PE:
	This is to avoid computing PE again and again in repeated encoder blocks.
	
	Arguments:
	x: input tensor of shape (batch_size, text_len, input_dim)
	
	Output:
	pe: tensor of shape (text_len, input_dim)
	pe[position, 2i] = sin( position * 10000^(- 2i / input_dim) )
	pe[position, 2i+1] = cos( position * 10000^(- 2i / input_dim) )
	"""
	_, text_len, input_dim = x.size()
	
	position = torch.arange(text_len, dtype = torch.float, device = x.device) ## shape (text_len, )
	
	div_term = torch.arange(0, input_dim, 2, dtype = torch.float, device = x.device) ##shape (input_dim//2, )
	div_term = torch.pow(10000, - div_term/input_dim)
	
	## Compute angles: tensor of shape (text_len, input_dim //2) as the outer product of position and div_term
	## angles[position, i] = position * 10000^(- 2i / input_dim)
	angles = torch.ger(position, div_term)
	
	## Interweave sin(angles) and cos(angles)
	## shape (text_len, input_dim)
	pe = torch.stack( (torch.sin(angles), torch.cos(angles)), dim = 2).view(text_len, input_dim) 
	return pe