104 lines
3.5 KiB
Python
104 lines
3.5 KiB
Python
import jax
|
|
import jax.numpy as jnp
|
|
from jax import lax
|
|
from jax import nn
|
|
|
|
|
|
@jax.jit
|
|
def sigmoid(x):
|
|
return 1 / (1 + lax.exp(-x))
|
|
|
|
|
|
def LSTMCell(hidden_size, w_init=nn.initializers.glorot_normal(), b_init=nn.initializers.normal()):
|
|
def init_fun(rng, input_size):
|
|
k1, k2, k3 = jax.random.split(rng, 3)
|
|
weight_ih = w_init(k1, (input_size, 4 * hidden_size))
|
|
weight_hh = w_init(k2, (hidden_size, 4 * hidden_size))
|
|
bias = b_init(k3, (4 * hidden_size,))
|
|
return hidden_size, (weight_ih, weight_hh, bias)
|
|
|
|
@jax.jit
|
|
def apply_fun(params, inputs, states):
|
|
(hx, cx) = states
|
|
weight_ih, weight_hh, bias = params
|
|
gates = lax.dot(inputs, weight_ih) + lax.dot(hx, weight_hh) + bias
|
|
|
|
in_gate = sigmoid(gates[:, :16])
|
|
forget_gate = sigmoid(gates[:, 16:32])
|
|
cell_gate = lax.tanh(gates[:, 32:48])
|
|
out_gate = sigmoid(gates[:, 48:])
|
|
|
|
cy = lax.mul(forget_gate, cx) + lax.mul(in_gate, cell_gate)
|
|
hy = lax.mul(out_gate, lax.tanh(cy))
|
|
|
|
return (hy, cy)
|
|
|
|
return init_fun, apply_fun
|
|
|
|
|
|
def LSTMLayer(hidden_size, w_init=nn.initializers.glorot_normal(), b_init=nn.initializers.normal()):
|
|
cell_init, cell_fun = LSTMCell(hidden_size, w_init=w_init, b_init=b_init)
|
|
|
|
@jax.jit
|
|
def apply_fun(params, inputs, states):
|
|
output = []
|
|
for input_data in inputs:
|
|
states = cell_fun(params, input_data, states)
|
|
output.append(states[0])
|
|
return jnp.stack(output), states
|
|
|
|
return cell_init, apply_fun
|
|
|
|
|
|
def StackedLSTM(hidden_size, num_layer, w_init=nn.initializers.glorot_normal(), b_init=nn.initializers.normal()):
|
|
layer_inits = []
|
|
layer_funs = []
|
|
for _ in range(num_layer):
|
|
layer_init, layer_fun = LSTMLayer(hidden_size, w_init=w_init, b_init=b_init)
|
|
layer_inits.append(layer_init)
|
|
layer_funs.append(layer_fun)
|
|
del layer_init
|
|
del layer_fun
|
|
|
|
def init_fun(rng, input_size):
|
|
params = []
|
|
output_size = input_size
|
|
for init in layer_inits:
|
|
output_size, layer_params = init(rng, output_size)
|
|
params.append(layer_params)
|
|
return output_size, tuple(params)
|
|
|
|
@jax.jit
|
|
def apply_fun(params, inputs, states):
|
|
output = inputs
|
|
output_states = []
|
|
for layer_id in range(num_layer):
|
|
output, out_state = layer_funs[layer_id](params[layer_id], output, states[layer_id])
|
|
output_states.append(out_state)
|
|
return jnp.stack(output), output_states
|
|
|
|
return init_fun, apply_fun
|
|
|
|
|
|
def StackedLSTMModel(input_size, hidden_size=-1, output_size=-1, num_layer=3,
|
|
w_init=nn.initializers.glorot_normal(), b_init=nn.initializers.normal()):
|
|
hidden_size = hidden_size if hidden_size > 0 else input_size * 2
|
|
output_size = output_size if output_size > 0 else input_size
|
|
stacked_init, stacked_fun = StackedLSTM(hidden_size, num_layer, w_init=w_init, b_init=b_init)
|
|
|
|
def init_fun(rng, input_size):
|
|
stacked_output_size, stacked_params = stacked_init(rng, input_size)
|
|
|
|
k1, k2 = jax.random.split(rng)
|
|
weight = w_init(k1, (stacked_output_size, output_size))
|
|
bias = b_init(k2, (output_size,))
|
|
|
|
return output_size, ((weight, bias), stacked_params)
|
|
|
|
@jax.jit
|
|
def apply_fun(params, inputs, states):
|
|
(weight, bias), stacked_params = params
|
|
output, new_states = stacked_fun(stacked_params, inputs, states)
|
|
return lax.dot(output[-1], weight) + bias, new_states
|
|
|
|
return init_fun, apply_fun
|