@jtu.skip_on_flag('jax_skip_slow_tests', True)
def testScanRnn(self):
    r = npr.RandomState(0)
    n_in = 4
    n_hid = 2
    n_out = 1
    length = 3
    W_trans = r.randn(n_hid, (n_hid + n_in)).astype(jnp.float_)
    W_out = r.randn(n_out, (n_hid + n_in)).astype(jnp.float_)
    params = (W_trans, W_out)
    inputs = r.randn(length, n_in).astype(jnp.float_)
    targets = r.randn(length, n_out).astype(jnp.float_)

    def step(params, state, input):
        (W_trans, W_out) = params
        stacked = jnp.concatenate([state, input])
        output = jnp.tanh(jnp.dot(W_out, stacked))
        next_state = jnp.tanh(jnp.dot(W_trans, stacked))
        return (next_state, output)

    def rnn(params, inputs):
        init_state = jnp.zeros(n_hid)
        (_, outputs) = lax.scan(partial(step, params), init_state, inputs)
        return outputs

    def loss(params, inputs, targets):
        predictions = rnn(params, inputs)
        return jnp.sum(((predictions - targets) ** 2))
    loss(params, inputs, targets)
    api.jvp((lambda params: loss(params, inputs, targets)), (params,), (params,))
    jtu.check_grads(loss, (params, inputs, targets), order=2, modes=['fwd'], rtol={np.float32: 0.02, np.float64: 1e-06})
    (_, expected) = api.jvp(loss, (params, inputs, targets), (params, inputs, targets))
    (_, linfun) = api.linearize(loss, params, inputs, targets)
    ans = linfun(params, inputs, targets)
    self.assertAllClose(ans, expected, check_dtypes=False)
    api.grad(loss)(params, inputs, targets)
    jtu.check_grads(loss, (params, inputs, targets), order=2, rtol=0.02)
    batch_size = 7
    batched_inputs = r.randn(batch_size, length, n_in).astype(jnp.float_)
    batched_targets = r.randn(batch_size, length, n_out).astype(jnp.float_)
    batched_loss = api.vmap((lambda x, y: loss(params, x, y)))
    losses = batched_loss(batched_inputs, batched_targets)
    expected = np.stack(list(map((lambda x, y: loss(params, x, y)), batched_inputs, batched_targets)))
    self.assertAllClose(losses, expected, check_dtypes=False, rtol=0.01)
@jtu.skip_on_flag('jax_skip_slow_tests', True)
def testScanRnn(self):
    r = npr.RandomState(0)
    n_in = 4
    n_hid = 2
    n_out = 1
    length = 3
    W_trans = r.randn(n_hid, (n_hid + n_in)).astype(jnp.float_)
    W_out = r.randn(n_out, (n_hid + n_in)).astype(jnp.float_)
    params = (W_trans, W_out)
    inputs = r.randn(length, n_in).astype(jnp.float_)
    targets = r.randn(length, n_out).astype(jnp.float_)

    def step(params, state, input):
        (W_trans, W_out) = params
        stacked = jnp.concatenate([state, input])
        output = jnp.tanh(jnp.dot(W_out, stacked))
        next_state = jnp.tanh(jnp.dot(W_trans, stacked))
        return (next_state, output)

    def rnn(params, inputs):
        init_state = jnp.zeros(n_hid)
        (_, outputs) = lax.scan(partial(step, params), init_state, inputs)
        return outputs

    def loss(params, inputs, targets):
        predictions = rnn(params, inputs)
        return jnp.sum(((predictions - targets) ** 2))
    loss(params, inputs, targets)
    api.jvp((lambda params: loss(params, inputs, targets)), (params,), (params,))
    jtu.check_grads(loss, (params, inputs, targets), order=2, modes=['fwd'], rtol={np.float32: 0.02, np.float64: 1e-06})
    (_, expected) = api.jvp(loss, (params, inputs, targets), (params, inputs, targets))
    (_, linfun) = api.linearize(loss, params, inputs, targets)
    ans = linfun(params, inputs, targets)
    self.assertAllClose(ans, expected, check_dtypes=False)
    api.grad(loss)(params, inputs, targets)
    jtu.check_grads(loss, (params, inputs, targets), order=2, rtol=0.02)
    batch_size = 7
    batched_inputs = r.randn(batch_size, length, n_in).astype(jnp.float_)
    batched_targets = r.randn(batch_size, length, n_out).astype(jnp.float_)
    batched_loss = api.vmap((lambda x, y: loss(params, x, y)))
    losses = batched_loss(batched_inputs, batched_targets)
    expected = np.stack(list(map((lambda x, y: loss(params, x, y)), batched_inputs, batched_targets)))
    self.assertAllClose(losses, expected, check_dtypes=False, rtol=0.01)