| jax | numpy | tensorflow | torch | |
|---|---|---|---|---|
| activation |  |
 |
 |
 |
| creation |  |
 |
 |
|
| device |  |
 |
 |
 |
| dtype |  |
 |
 |
 |
| elementwise |  |
 |
 |
 |
| general |  |
 |
 |
 |
| gradients |  |
 |
 |
 |
| layers |  |
 |
 |
 |
| linalg |  |
 |
 |
 |
| losses |  |
 |
 |
 |
| manipulation |  |
 |
 |
 |
| meta |  |
 |
 |
 |
| nest |  |
 |
 |
 |
| norms |  |
 |
 |
 |
| random |  |
 |
 |
 |
| searching |  |
 |
 |
 |
| set |  |
 |
 |
 |
| sorting |  |
 |
 |
 |
| statistical |  |
 |
 |
 |
| uitility |  |
 |
 |
 |