def test_embed(self):
    rng = dict(params=random.PRNGKey(0))
    x = jnp.arange(4)[None]
    dummy_embedding = jnp.broadcast_to(jnp.arange(4)[(..., None)], (4, 3)).astype(jnp.float32)
    embed_module = nn.Embed(num_embeddings=4, features=3, embedding_init=(lambda rng, shape, dtype: dummy_embedding))
    (y, initial_params) = embed_module.init_with_output(rng, x)
    np.testing.assert_allclose(y, dummy_embedding[None])
    z = embed_module.apply(initial_params, jnp.ones((3,)), method=embed_module.attend)
    np.testing.assert_allclose(z, (3.0 * jnp.arange(4)))
def test_embed(self):
    rng = dict(params=random.PRNGKey(0))
    x = jnp.arange(4)[None]
    dummy_embedding = jnp.broadcast_to(jnp.arange(4)[(..., None)], (4, 3)).astype(jnp.float32)
    embed_module = nn.Embed(num_embeddings=4, features=3, embedding_init=(lambda rng, shape, dtype: dummy_embedding))
    (y, initial_params) = embed_module.init_with_output(rng, x)
    np.testing.assert_allclose(y, dummy_embedding[None])
    z = embed_module.apply(initial_params, jnp.ones((3,)), method=embed_module.attend)
    np.testing.assert_allclose(z, (3.0 * jnp.arange(4)))