Hello,
I am trying to implement Simple Recurrent Unit (SRU) 21. The core idea of SRU lies in Equation (3)-(7) and my naive implementation (i.e., without any optimization) for bi-SRU is below:
def SRUStep(ft, rt, xt_tilde, ctm1):
	"""
	"""
	tanh = nn.Tanh()
	c_t = ft * ctm1 + (1.0 - ft) * xt_tilde
	h_t = rt * tanh(c_t) + (1.0 - rt) * xt_tilde
	return h_t, c_t

class SRU(nn.Module):
	"""
	implementation of simple recurrent unit
	"""
	def __init__(self, args):
		"""
		"""
		super(SRU, self).__init__()
		self.n_in = args.dim_w
		self.n_out = args.dim_h
		self.W = Parameter(torch.FloatTensor(np.zeros((self.n_in, 3 * self.n_out))))
		self.bf = Parameter(torch.FloatTensor(np.zeros(self.n_out)))
		self.br = Parameter(torch.FloatTensor(np.zeros(self.n_out)))

		# recurrent dropout
		self.Dropout_recurrent = nn.Dropout(0.5)
		# non-recurrent dropout
		self.Dropout = nn.Dropout(0.3)

		self.sigmoid = nn.Sigmoid()
		self.tanh = nn.Tanh()
		self.bs = args.bs
		self.n_steps = args.sent_len 
		self.init_weight()

	def forward(self, x):
		"""
		x shape: (bs, max_len, dim_x)
		"""
		x = self.Dropout(x)
		# shape: (bs, max_len, 3 * dim_h)
		Wx = torch.bmm(x, self.W.repeat(self.bs, 1, 1))
		# using different dropout masks at different time steps
		Wx = self.Dropout_recurrent(Wx)
		x_tilde = Wx[:, :, :self.n_out]
		f = self.sigmoid(Wx[:, :, self.n_out:2*self.n_out] + self.bf)
		r = self.sigmoid(Wx[:, :, 2*self.n_out:] + self.br)
		ct = Variable(x.data.new(self.bs, self.n_out).zero_())
		H_fwd, H_bwd = [], []
		for i in range(self.n_steps):
			ft = f[:, i]
			rt = r[:, i]
			xt_tilde = x_tilde[:, i]
			ht, ct = SRUStep(ft, rt, xt_tilde, ct)
			H_fwd.append(ht)
		ct = Variable(x.data.new(self.bs, self.n_out).zero_())
		for i in range(self.n_steps-1, -1, -1):
			ft = f[:, i]
			rt = r[:, i]
			xt_tilde = x_tilde[:, i]
			ht, ct = SRUStep(ft, rt, xt_tilde, ct)
			H_bwd.append(ht)
		H_fwd = torch.stack(H_fwd).permute(1, 0, 2)
		H_bwd = torch.stack(H_bwd).permute(1, 0, 2)
		H = torch.cat([H_fwd, H_bwd], dim=2)
		return H

	def init_weight(self):
		"""
		"""
		val_range = (3.0 / self.n_in) ** 0.5
		# sample value from a uniform distribution
		self.W.data.uniform_(-val_range, val_range)

I wonder if my implementation is correct in Pytorch