PSO-JAX

An accelerated Particle Swarm Optimization tool that uses JAX key component. Here are the main benefits:

Runs on Python.
Supports multi-core CPUs, GPUs, and TPUs.
Compatible with Colab environment.

Getting Started

See this example

References

BHI Research
JAX
PySwarms is used in the examples to compare results.