LiibanMo / scikit-jax
Your favourite classical machine learning algos on the GPU/TPU
☆20Updated 3 months ago
Alternatives and similar repositories for scikit-jax:
Users that are interested in scikit-jax are comparing it to the libraries listed below
- ☆27Updated 9 months ago
- nanoGPT using Equinox☆13Updated 2 years ago
- Jax like function transformation engine but micro, microjax☆30Updated 6 months ago
- Generative cellular automaton-like learning environments for RL.☆19Updated 2 months ago
- Graph neural networks in JAX.☆67Updated 10 months ago
- A simple example of VAEs with KANs☆12Updated 11 months ago
- An implementation of ESM2 in Equinox+JAX☆25Updated 2 months ago
- gzip Predicts Data-dependent Scaling Laws☆34Updated 10 months ago
- Neural Networks for JAX☆84Updated 7 months ago
- This is a port of Mistral-7B model in JAX☆32Updated 9 months ago
- Proof-of-concept of global switching between numpy/jax/pytorch in a library.☆18Updated 10 months ago
- ☆102Updated this week
- Source code for the paper "Positional Attention: Out-of-Distribution Generalization and Expressivity for Neural Algorithmic Reasoning"☆14Updated 2 months ago
- JAX/Flax implementation of the Hyena Hierarchy☆34Updated last year
- Agent framework for constructing language model agents and training on constructive tasks.☆72Updated this week
- An annotated implementation of the Hyena Hierarchy paper☆32Updated last year
- ☆47Updated last month
- ☆43Updated last month
- Fine-grained, dynamic control of neural network topology in JAX.☆21Updated last year
- ☆20Updated last year
- Scalable and Stable Parallelization of Nonlinear RNNS☆15Updated 2 months ago
- ☆21Updated 3 weeks ago
- Code for minimum-entropy coupling.☆31Updated 9 months ago
- ☆60Updated 3 years ago
- ☆38Updated 9 months ago
- Implementation of Gradient Agreement Filtering, from Chaubard et al. of Stanford, but for single machine microbatches, in Pytorch☆24Updated 3 months ago
- Triton Implementation of HyperAttention Algorithm☆47Updated last year
- Implementation of the Pairformer model used in AlphaFold 3☆12Updated this week
- Simple Scalable Discrete Diffusion for text in PyTorch☆33Updated 6 months ago
- Because we don't want a jupyter notebook mess...☆62Updated last month