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.
Built Distribution
File details
Details for the file gateloop-0.0.2-py3-none-any.whl.
File metadata
- Download URL: gateloop-0.0.2-py3-none-any.whl
- Upload date:
- Size: 3.2 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/4.0.2 CPython/3.10.3
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 | 7ca1a95bb558dcbc409754c5fb216b117d142a68c3aaf8aad53673a5f0ac1690 |
|
| MD5 | e48ac4ccba396f4a09a5bb0341ce295e |
|
| BLAKE2b-256 | 80d9fec54186bf207e2945a1410614b628e7d2cc0ba5fa3745f912a66872e823 |
