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

Name		Name	Last commit message	Last commit date
Latest commit History 12 Commits
weights		weights
.gitignore		.gitignore
LICENSE		LICENSE
README.md		README.md
data.py		data.py
nn.py		nn.py
plot.py		plot.py
train.py		train.py