No project description provided
Project description
GateLoop
Implementation of arxiv.org/abs/2311.01927 in JAX based on lucidrains/gateloop-transformer.
Installation
pip install gateloop
Usage
from gateloop import GateLoopTransformer
key = jax.random.PRNGKey(0)
model = GateLoopTransformer(
num_tokens = 1000,
dim = 512,
depth = 12
)
params = model.init(key, jnp.ones((1, 10), jnp.int32))
ids = jax.random.randint(key, (1, 10), 0, 1000)
logits = model.apply(params, ids)
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distributions
No source distribution files available for this release.See tutorial on generating distribution archives.
