Alternatives and similar repositories for jax_generative_models

Users that are interested in jax_generative_models are comparing it to the libraries listed below

• kvfrans / lmpo
  ☆134Updated last month
• clement-bonnet / lpn
  Latent Program Network (from the "Searching Latent Program Spaces" paper)
  ☆107Updated 2 months ago
• sdan / nanoEBM
  minimal Energy-based transformer
  ☆42Updated last month
• nmboffi / jax-interpolants
  Definitive implementation of the stochastic interpolant framework for generative modeling in jax.
  ☆40Updated 6 months ago
• kvfrans / splus
  ☆123Updated 7 months ago
• vmicheli / delta-iris
  Efficient World Models with Context-Aware Tokenization. ICML 2024
  ☆115Updated last year
• maxencefaldor / omni-epic
  OMNI-EPIC: Open-endedness via Models of human Notions of Interestingness with Environments Programmed in Code (ICLR 2025).
  ☆73Updated last year
• iliao2345 / CompressARC
  ☆214Updated 3 weeks ago
• p-doom / jasmine
  A simple, performant and scalable JAX-based world modeling codebase.
  ☆123Updated 2 weeks ago
• garymm / earl
  Reinforcement learning with Equinox
  ☆20Updated 10 months ago
• yilundu / ired_code_release
  ☆83Updated last year
• kvfrans / jax-diffusion-transformer
  Implementation of Diffusion Transformer (DiT) in JAX
  ☆305Updated last year
• smsharma / minified-generative-models
  Bare-bones implementations of some generative models in Jax: diffusion, normalizing flows, consistency models, flow matching, (beta)-VAEs…
  ☆141Updated 2 years ago
• lucidrains / evolutionary-policy-optimization
  Pytorch implementation of Evolutionary Policy Optimization, from Wang et al. of the Robotics Institute at Carnegie Mellon University
  ☆103Updated 4 months ago
• okarthikb / state-space-models
  ☆27Updated last year
• lucidrains / improving-transformers-world-model-for-rl
  Implementation of the new SOTA for model based RL, from the paper "Improving Transformer World Models for Data-Efficient RL", in Pytorch
  ☆151Updated 8 months ago
• facebookresearch / minimax
  Efficient baselines for autocurricula in JAX.
  ☆206Updated last year
• model-architectures / GRAPE
  Official implementation of GRAPE: Group Representational Position Encoding (https://arxiv.org/abs/2512.07805)
  ☆74Updated 3 weeks ago
• maxencefaldor / cax
  Cellular Automata Accelerated in JAX (Oral at ICLR 2025)
  ☆239Updated 2 months ago
• smearle / autoverse
  Generative cellular automaton-like learning environments for RL.
  ☆20Updated last year
• facebookresearch / adjoint_sampling
  code for "Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching"
  ☆129Updated 6 months ago
• idiap / sigma-gpt
  σ-GPT: A New Approach to Autoregressive Models
  ☆70Updated last year
• kvfrans / jax-flow
  Flow-matching algorithms in JAX
  ☆114Updated last year
• lucidrains / transformer-directed-evolution
  Explorations into whether a transformer with RL can direct a genetic algorithm to converge faster
  ☆71Updated 8 months ago
• NX-AI / xlstm-jax
  Official JAX implementation of xLSTM including fast and efficient training and inference code. 7B model available at https://huggingface.…
  ☆105Updated last year
• arcprize / hierarchical-reasoning-model-analysis
  ☆167Updated 5 months ago
• facebookresearch / lagrangian-ot
  Neural Optimal Transport with Lagrangian Costs
  ☆61Updated 8 months ago
• LiibanMo / scikit-jax
  Your favourite classical machine learning algos on the GPU/TPU
  ☆21Updated last month
• imbue-ai / carbs
  Cost aware hyperparameter tuning algorithm
  ☆177Updated last year
• nebius / kvax
  A FlashAttention implementation for JAX with support for efficient document mask computation and context parallelism.
  ☆157Updated 2 months ago