Metadata-Version: 2.1
Name: jaxkan
Version: 0.1.9
Summary: A JAX implementation of Kolmogorov-Arnold Networks
Author-email: Spyros Rigas <rigassp@gmail.com>, Michalis Papachristou <mixpap@phys.uoa.gr>
License: MIT License
        
        Copyright (c) 2024 Spyros Rigas, Michalis Papachristou
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
        
Project-URL: Repository, https://github.com/srigas/jaxkan
Project-URL: Issues, https://github.com/srigas/jaxkan/issues
Keywords: JAX,NNX,KANs,Kolmogorov-Arnold,PIKAN
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: License :: OSI Approved :: MIT License
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy==2.1.3
Requires-Dist: scipy==1.14.1
Requires-Dist: jax[cpu]==0.4.35
Requires-Dist: optax==0.2.4
Requires-Dist: flax==0.10.2
Requires-Dist: pytest==8.3.4
Provides-Extra: gpu
Requires-Dist: jax[cuda12]==0.4.35; extra == "gpu"
Provides-Extra: doc
Requires-Dist: jupyterlab==4.3.0; extra == "doc"
Requires-Dist: matplotlib==3.9.2; extra == "doc"
Requires-Dist: scikit-learn==1.5.2; extra == "doc"
Requires-Dist: sphinx; extra == "doc"
Requires-Dist: sphinx-rtd-theme; extra == "doc"
Requires-Dist: sphinx-autobuild; extra == "doc"
Requires-Dist: myst-parser; extra == "doc"
Requires-Dist: sphinx_design; extra == "doc"
Requires-Dist: sphinx-autodoc-typehints; extra == "doc"
Requires-Dist: nbsphinx; extra == "doc"
Requires-Dist: nbsphinx-link; extra == "doc"

# jaxKAN

A JAX implementation of the original Kolmogorov-Arnold Networks (KANs), using the Flax and Optax frameworks for neural networks and optimization, respectively. Our adaptation is based on the original [pykan](https://github.com/KindXiaoming/pykan), however we also included a built-in grid extension routine, which does not simply perform an adaptation of the grid based on the inputs, but also extends its size.


## Installation

`jaxKAN` is available as a PyPI package. For installation, simply run

```
pip3 install jaxkan
```

The default installation requires `jax[cpu]`, but there is also a `gpu` version which will install `jax[cuda12]` as a dependency.


## Why not more efficient?

Despite their overall potential in the Deep Learning field, the authors of KANs emphasized their performance when it comes to scientific computing, in tasks such as Symbolic Regression or solving PDEs. This is why we put emphasis on preserving their original form, albeit less computationally efficient, as it allows the user to utilize the full regularization terms presented in the [arXiv pre-print](https://arxiv.org/abs/2404.19756) and not the "mock" regularization terms presented, for instance, in the [efficient-kan](https://github.com/Blealtan/efficient-kan/tree/master) implementation.


## Citation

If you utilized `jaxKAN` for your own academic work, please consider using the following citation, which is the paper introducing the framework:

```
@article{10763509,
      author = {Rigas, Spyros and Papachristou, Michalis and Papadopoulos, Theofilos and Anagnostopoulos, Fotios and Alexandridis, Georgios},
      journal = {IEEE Access}, 
      title = {Adaptive Training of Grid-Dependent Physics-Informed Kolmogorov-Arnold Networks}, 
      year = {2024},
      volume = {12},
      pages = {176982-176998},
      doi = {10.1109/ACCESS.2024.3504962}
}
```
