Neural Operators in JAX
Project description
neojax
[!WARNING] Work in Progress: This project is in its early stages of development and is currently not functional. It is being shared for development and architectural purposes only.
This is neojax (Neural Operators in jax), a port of the popular neural operators pytorch library to jax. It is built on top of jax and equinox and provides the same API as the original pytorch library.
Installation
Install the python package via pypi
pip3 install neojax-operators
Quickstart
neojax exposes the same API as neuraloperators and can therefore be used as a drop-in replacement:
from neojax.models import FNO
operator = FNO(n_modes=(64, 64),
hidden_channels=64,
in_channels=2,
out_channels=1)
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file neojax_operators-0.0.2.tar.gz.
File metadata
- Download URL: neojax_operators-0.0.2.tar.gz
- Upload date:
- Size: 604.2 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.2.0 CPython/3.11.15
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
a6b6a2b6a70aec7705948ec9c51cc31eaf620dfb35a1afcc261b6528a7f77e9f
|
|
| MD5 |
720fc9d2b4353a9b546172ec7833b3e2
|
|
| BLAKE2b-256 |
105d13d45d86de7848d3ad04019d28b6934e34cb5bcadbac49616a578af12c66
|
File details
Details for the file neojax_operators-0.0.2-py3-none-any.whl.
File metadata
- Download URL: neojax_operators-0.0.2-py3-none-any.whl
- Upload date:
- Size: 13.4 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.2.0 CPython/3.11.15
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
a728870852e9d720d2526e334488b800d84fba80d1760a6b15d7af0052024a41
|
|
| MD5 |
476b2996bd27512f5879268aef5fcda3
|
|
| BLAKE2b-256 |
284c5bea6a8f16719170f429ce3c6b0dcb6f427a00cb860df56c215efbddaa6b
|
