vrp-attention-jax

Reimplementation of the paper "Attention, Learn to Solve Routing Problems!" in jax/flax.

MIT License

Stars
7
View Code on GitHub View on X
Ecosystems: Python

JAX + Attention Learn To Solve Routing Problems

Reimplementation of the paper Attention, Learn to Solve Routing Problems! using Jax and Flax. Fully suitable for multiple GPU and TPU execution.

Huge thanks to previous implementations in PyTorch and TensorFlow.

Reproducing the original experiments

Download the evaluation data.

Related Projects
tensorf-jax

Unofficial implementation of Tensorial Radiance Fields (Chen & Xu ‘22)

17 Apr 2022 39
mlp-gpt-jax

A GPT, made only of MLPs, in Jax

21 May 2021 55
fidjax

Frechet Inception Distance in JAX

15 Jul 2024 1
flax

Flax is a neural network library for JAX that is designed for flexibility.

10 Jan 2020 5,584
jaxlie

Rigid transforms + Lie groups in JAX

28 Nov 2020 220
pips-jax

JAX port of Persistent Independent Particles

20 Jan 2023 36
minGPT-flax

GPT implementation in Flax

30 Sep 2021 16
lpips-j

Minimal JAX/Flax port of `lpips` supporting `vgg16`, with pre-trained weights stored in the 🤗 Hug...

08 Jul 2022 14
attention-is-all-you-need

Pytorch implementation of the reknowned "Attention Is All You Need" paper - NeurIPS 2017

21 Jan 2024 3
rlax

18 Feb 2020 1,193
n-grammer-flax

Implementation of N-Grammer in Flax

04 Aug 2022 14
vqgan-jax

JAX implementation of VQGAN

28 Jun 2021 90
vit-vqgan

JAX implementation ViT-VQGAN

22 Jun 2022 77
equinox

Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/

29 Jul 2021 1,846
flash-attention-jax

Implementation of Flash Attention in Jax

12 Jul 2022 189