neojax-operators 0.0.1__tar.gz

neojax_operators-0.0.1/.gitignore

@@ -0,0 +1,2 @@
+.vscode/*
+

neojax_operators-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Paul Gekeler
+
+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.

neojax_operators-0.0.1/PKG-INFO

@@ -0,0 +1,37 @@
+Metadata-Version: 2.4
+Name: neojax-operators
+Version: 0.0.1
+Summary: Neural Operators in JAX
+Project-URL: Homepage, https://github.com/paulgekeler/neojax
+Author-email: Paul Gekeler <pgekeler@gmail.com>
+License: MIT
+License-File: LICENSE
+Keywords: deep-learning,fno,jax,neural-operators,scientific-computing
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Requires-Dist: equinox
+Requires-Dist: jax
+Requires-Dist: jaxlib
+Description-Content-Type: text/markdown
+
+This is neojax-operators (**Ne**ural **O**perators in **jax**), a port of the popular [neural operators](https://github.com/neuraloperator/neuraloperator) pytorch library to jax. It is built on top of [jax](https://github.com/jax-ml/jax) and [equinox](https://github.com/patrick-kidger/equinox) and provides the same API as the original pytorch library.
+
+#### Installation
+Install the python package via pypi
+```bash
+pip3 install neojax-operators
+```
+
+#### Quickstart
+neojax-operators exposes the same API as neuraloperators and can therefore be used as a drop-in replacement:
+```python
+from neojax.models import FNO
+operator = FNO(n_modes=(64, 64),
+               hidden_channels=64,
+               in_channels=2,
+               out_channels=1)
+```

neojax_operators-0.0.1/README.md

@@ -0,0 +1,17 @@
+This is neojax-operators (**Ne**ural **O**perators in **jax**), a port of the popular [neural operators](https://github.com/neuraloperator/neuraloperator) pytorch library to jax. It is built on top of [jax](https://github.com/jax-ml/jax) and [equinox](https://github.com/patrick-kidger/equinox) and provides the same API as the original pytorch library.
+
+#### Installation
+Install the python package via pypi
+```bash
+pip3 install neojax-operators
+```
+
+#### Quickstart
+neojax-operators exposes the same API as neuraloperators and can therefore be used as a drop-in replacement:
+```python
+from neojax.models import FNO
+operator = FNO(n_modes=(64, 64),
+               hidden_channels=64,
+               in_channels=2,
+               out_channels=1)
+```

neojax_operators-0.0.1/neojax/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.0.1"

neojax_operators-0.0.1/neojax/layers/attention_kernel_integral.py

@@ -0,0 +1,300 @@
+from typing import Callable, Any
+import math
+import jax
+from jax import Array
+import jax.numpy as jnp
+import equinox as eqx
+from .instance_norm1d import InstanceNorm1d
+
+
+class AttentionKernelIntegral(eqx.Module):
+    """Kernel integral transform with attention
+
+    Computes \\int_{Omega} k(x, y) * f(y) dy,
+    where:
+          K(x, y) = \\sum_{c=1}^d q_c(x) * k_c(y), q(x) = [q_1(x); ...; q_d(x)], k(y) = [k_1(y); ...; k_d(y)]
+          f(y) = v(y)
+    More specifically, this module supports using just one input function (self-attention) or
+    two input functions (cross-attention) to compute the kernel integral transform.
+
+    1. Self-attention:
+        input function u(.), sampling grid D_x = {x_i}_{i=1}^N
+        query function: q(x_i) = u(x_i) W_q
+        key function: k(x_i) = u(x_i) W_k
+        value function: v(x_i) = u(x_i) W_v
+
+    2. Cross-attention:
+        first input function u_qry(.), sampling grid D_x = {x_i}_{i=1}^N
+        second input function u_src(.), sampling grid D_y = {y_j}_{j=1}^M, D_y can be different from D_x
+        query function: q(x_i) = u_qry(x_i) W_q
+        key function: k(y_j) = u_src(y_j) W_k
+        value function: v(y_j) = u_src(y_j) W_v
+
+    Self-attention can be considered as a special case of cross-attention, where u = u_qry = u_src and D_x = D_y.
+
+    The kernel integral transform will be numerically computed as:
+        \\int_{Omega} k(x, y) * f(y) dy \\appox \\sum_{j=1}^M * k(x, y_j) * f(y_j) * w(y_j)
+    For uniform quadrature, the weights w(y_j) = 1/M.
+    For non-uniform quadrature, the weights w(y_j) is specified as an input to the forward function.
+
+    Parameters
+    ----------
+    in_channels : int
+        Number of input channels
+    out_channels : int
+        Number of output channels
+    n_heads : int
+        Number of attention heads in multi-head attention
+    head_n_channels : int
+        Dimension of each attention head, determines how many function bases to use for the kernel
+        k(x, y) = \\sum_{c=1}^d \\q_c(x) * \\k_c(y), head_n_channels controls the d
+    pos_dim : int
+        Dimension of the domain, determines the dimension of coordinates
+    project_query : bool, optional
+        Whether to project the query function with pointwise linear layer
+        (this is sometimes not needed when using cross-attention), by default True
+    """
+
+    def __init__(
+        self,
+        key: Array,
+        in_channels: int,
+        out_channels: int,
+        n_heads: int,
+        head_n_channels: int,
+        project_query: bool = True,
+    ) -> None:
+        super().__init__()
+        l1_key, l2_key, l3_key, l4_key = jax.random.split(key, 4)
+        self.n_heads = n_heads
+        self.head_n_channels = head_n_channels
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        self.project_query = project_query
+        if project_query:
+            self.to_q = eqx.nn.Linear(
+                in_channels, head_n_channels * n_heads, use_bias=False, key=l1_key
+            )
+        else:
+            self.to_q = eqx.nn.Identity()
+
+        self.to_k = eqx.nn.Linear(
+            in_channels, head_n_channels * n_heads, use_bias=False, key=l2_key
+        )
+
+        self.k_norm = InstanceNorm1d(head_n_channels)
+        self.v_norm = InstanceNorm1d(head_n_channels)
+
+        self.to_v = eqx.nn.Linear(
+            in_channels, head_n_channels * n_heads, use_bias=False, key=l3_key
+        )
+
+        self.to_out = (
+            eqx.nn.Linear(head_n_channels * n_heads, out_channels, key=l4_key)
+            if head_n_channels * n_heads != out_channels
+            else eqx.nn.Identity()
+        )
+
+        self.init_gain = 1 / math.sqrt(head_n_channels)
+        self.diagonal_weight = self.init_gain
+        self.initialize_qkv_weights()
+
+    def init_weight(
+        self, weight: Array, init_fn: Callable[[Array, Any], Array]
+    ) -> None:
+        """
+        Initialization for the projection matrix
+        basically initialize the weights for each heads with predefined initialization function and gain,
+        to add the diagonal bias, it requires input channels = head_n_channels
+        W = init_fn(W) + I * diagonal_weight
+
+        init_fn is xavier_uniform_ by default
+        """
+
+        for param in weight.parameters():
+            if param.ndim > 1:
+                for h in range(self.n_heads):
+                    init_fn(
+                        param[
+                            h * self.head_n_channels : (h + 1) * self.head_n_channels, :
+                        ],
+                        gain=self.init_gain,
+                    )
+                    if self.head_n_channels == self.in_channels:
+                        diagonal_bias = self.diagonal_weight * jnp.diag(
+                            jnp.ones(param.size(-1), dtype=jnp.float32)
+                        )
+                        param.data[
+                            h * self.head_n_channels : (h + 1) * self.head_n_channels, :
+                        ] += diagonal_bias
+
+    def initialize_qkv_weights(self) -> None:
+        """
+        Initialize the weights for q, k, v projection matrix with a small gain and add a diagonal bias,
+        this technique has been found useful for scale-sensitive problem that has not been normalized
+        see Table 8 in https://arxiv.org/pdf/2105.14995.pdf
+        """
+        init_fn = jax.nn.initializers.xavier_uniform()
+
+        if self.project_query:
+            self.init_weight(self.to_q, init_fn)
+        self.init_weight(self.to_k, init_fn)
+        self.init_weight(self.to_v, init_fn)
+
+    def normalize_wrt_domain(self, u: Array, norm_fn):
+        """
+        Normalize the input function with respect to the domain,
+         reshape the tensor to [batch_size*n_heads, num_grid_points, head_n_channels]
+        The second dimension is equal to the number of grid points that discretize the domain
+        """
+        # u: the input or transformed function
+        batch_size = u.shape[0]
+        u = u.view(batch_size * self.n_heads, -1, self.head_n_channels)
+        u = norm_fn(
+            u
+        )  # layer norm with channel dimension or instance norm with spatial dimension
+        return u.view(batch_size, self.n_heads, -1, self.head_n_channels)
+
+    def forward(
+        self,
+        u_src,
+        pos_src,
+        positional_embedding_module: eqx.nn.RotaryPositionalEmbedding = None,  # positional encoding module for encoding q/k
+        u_qry=None,
+        pos_qry=None,
+        weights=None,
+        associative=True,  # can be much faster if num_grid_points is larger than the channel number c
+        return_kernel=False,
+    ):
+        """
+        Computes kernel integral transform with attention
+
+        Parameters
+        ----------
+        u_src: input function (used to compute key and value in attention),
+                tensor of shape [batch_size, num_grid_points_src, channels]
+        pos_src: coordinate of the second source of function's sampling points y,
+                tensor of shape [batch_size, num_grid_points_src, pos_dim]
+        positional_embedding_module: positional embedding module for encoding query/key (q/k),
+                a torch.nn.Module
+        u_qry: query function,
+                tensor of shape [batch_size, num_grid_points_query, channels], if not provided, u_qry = u_src
+        pos_qry: coordinate of query points x,
+                tensor of shape [batch_size, num_grid_points_query, pos_dim], if not provided, pos_qry = pos_src
+        weights : quadrature weight w(y_j) for the kernel integral: u(x_i) = sum_{j} k(x_i, y_j) f(y_i) w(y_j),
+                tensor of shape [batch_size, num_grid_points_src], if not provided assume to be 1/num_grid_points_src
+        associative: if True, use associativity of matrix multiplication, first multiply K^T V, then multiply Q,
+                much faster when num_grid_points is larger than the channel number (which is usually the case)
+        return_kernel: if True, return the kernel matrix (for analyzing the kernel)
+
+        Output
+        ----------
+        u: Output function given on the query points x.
+        """
+
+        if u_qry is None:
+            u_qry = u_src  # go back to self attention
+            if pos_qry is not None:
+                raise ValueError(
+                    "Query coordinates are provided but query function is not provided"
+                )
+        else:
+            if pos_qry is None:
+                raise ValueError(
+                    "Query coordinates are required if query function is provided"
+                )
+
+        if return_kernel and associative:
+            raise ValueError("Cannot get kernel matrix when associative is set to True")
+
+        batch_size, num_grid_points = u_src.shape[
+            :2
+        ]  # batch size and number of grid points
+        pos_dim = pos_src.shape[-1]  # position dimension
+
+        q = self.to_q(u_qry)
+        k = self.to_k(u_src)
+        v = self.to_v(u_src)
+        q = (
+            q.view(batch_size, -1, self.n_heads, self.head_n_channels)
+            .permute(0, 2, 1, 3)
+            .contiguous()
+        )
+        k = (
+            k.view(batch_size, -1, self.n_heads, self.head_n_channels)
+            .permute(0, 2, 1, 3)
+            .contiguous()
+        )
+        v = (
+            v.view(batch_size, -1, self.n_heads, self.head_n_channels)
+            .permute(0, 2, 1, 3)
+            .contiguous()
+        )
+
+        k = self.normalize_wrt_domain(k, self.k_norm)
+        v = self.normalize_wrt_domain(v, self.v_norm)
+
+        if positional_embedding_module is not None:
+            if pos_dim == 2:
+                k_freqs_1 = positional_embedding_module.forward(pos_src[..., 0])
+                k_freqs_2 = positional_embedding_module.forward(pos_src[..., 1])
+                k_freqs_1 = k_freqs_1.unsqueeze(1).repeat([1, self.n_heads, 1, 1])
+                k_freqs_2 = k_freqs_2.unsqueeze(1).repeat([1, self.n_heads, 1, 1])
+
+                if pos_qry is None:
+                    q_freqs_1 = k_freqs_1
+                    q_freqs_2 = k_freqs_2
+                else:
+                    q_freqs_1 = positional_embedding_module.forward(pos_qry[..., 0])
+                    q_freqs_2 = positional_embedding_module.forward(pos_qry[..., 1])
+                    q_freqs_1 = q_freqs_1.unsqueeze(1).repeat([1, self.n_heads, 1, 1])
+                    q_freqs_2 = q_freqs_2.unsqueeze(1).repeat([1, self.n_heads, 1, 1])
+
+                q = positional_embedding_module.apply_2d_rotary_pos_emb(
+                    q, q_freqs_1, q_freqs_2
+                )
+                k = positional_embedding_module.apply_2d_rotary_pos_emb(
+                    k, k_freqs_1, k_freqs_2
+                )
+            elif pos_dim == 1:
+                k_freqs = positional_embedding_module.forward(pos_src[..., 0])
+                k_freqs = k_freqs.unsqueeze(1).repeat([batch_size, self.n_heads, 1, 1])
+
+                if pos_qry is None:
+                    q_freqs = k_freqs
+                else:
+                    q_freqs = positional_embedding_module.forward(pos_qry[..., 0])
+                    q_freqs = q_freqs.unsqueeze(1).repeat(
+                        [batch_size, self.n_heads, 1, 1]
+                    )
+
+                q = positional_embedding_module.apply_1d_rotary_pos_emb(q, q_freqs)
+                k = positional_embedding_module.apply_1d_rotary_pos_emb(k, k_freqs)
+            else:
+                raise ValueError(
+                    "Currently doesnt support relative embedding >= 3 dimensions"
+                )
+
+        if weights is not None:
+            weights = weights.view(batch_size, 1, num_grid_points, 1)
+        else:
+            weights = 1.0 / num_grid_points
+
+        if associative:
+            dots = jnp.matmul(k.transpose(-1, -2), v)
+            u = jnp.matmul(q, dots) * weights
+        else:
+            # this is more efficient when num_grid_points<<channels
+            kxy = jnp.matmul(q, k.transpose(-1, -2))
+            u = jnp.matmul(kxy, v) * weights
+
+        u = (
+            u.permute(0, 2, 1, 3)
+            .contiguous()
+            .view(batch_size, num_grid_points, self.n_heads * self.head_n_channels)
+        )
+        u = self.to_out(u)
+        if return_kernel:
+            return u, kxy
+        return u

neojax_operators-0.0.1/neojax/layers/instance_norm1d.py

@@ -0,0 +1,17 @@
+import jax
+from jax import Array
+import equinox as eqx
+
+
+class InstanceNorm1d(eqx.Module):
+    ln: eqx.nn.LayerNorm
+
+    def __init__(self, shape_per_channel: int) -> None:
+        # shape_per_channel is the length of the 1d signal
+        # we disable learnable affine parameters to match torch neuralop config
+        self.ln = eqx.nn.LayerNorm(
+            shape=(shape_per_channel,), use_weight=False, use_bias=False
+        )
+
+    def __call__(self, x: Array) -> Array:
+        return jax.vmap(self.ln)(x)

neojax_operators-0.0.1/neojax/utils.py

@@ -0,0 +1,234 @@
+from typing import Optional, Union
+from math import prod
+from pathlib import Path
+import jax
+import jax.numpy as jnp
+from jax import Array
+import equinox as eqx
+
+# Only import wandb and use if installed
+wandb_available = False
+try:
+    import wandb
+
+    wandb_available = True
+except ModuleNotFoundError:
+    wandb_available = False
+
+
+def count_model_params(model: eqx.Module) -> int:
+    """Returns the total number of parameters of an equinox model
+
+    Notes
+    -----
+    One complex number is counted as two parameters (we count real and imaginary parts)'
+    """
+    trainable_params = eqx.filter(model, eqx.is_inexact_array)
+    total_params = jax.vmap(count_array_params)(trainable_params)
+    return int(jnp.sum(total_params))
+
+
+def count_array_params(arr: Array, dims=None) -> int:
+    """Returns the number of parameters (elements) in a single array, optionally, along certain dimensions only
+
+    Parameters
+    ----------
+    array : Array
+    dims : int list or None, default is None
+        if not None, the dimensions to consider when counting the number of parameters (elements)
+
+    Notes
+    -----
+    One complex number is counted as two parameters (we count real and imaginary parts)'
+    """
+    if dims is None:
+        dims = list(arr.shape)
+    else:
+        dims = [arr.shape[d] for d in dims]
+    n_params = prod(dims)
+    if jnp.iscomplex(arr):
+        return 2 * n_params
+    return n_params
+
+
+def wandb_login(
+    api_key_file: str = "../config/wandb_api_key.txt", key: str = None
+) -> None:
+    if key is None:
+        key = get_wandb_api_key(api_key_file)
+
+    wandb.login(key=key)
+
+
+def set_wandb_api_key(api_key_file: str = "../config/wandb_api_key.txt") -> None:
+    import os
+
+    try:
+        os.environ["WANDB_API_KEY"]
+    except KeyError:
+        with open(api_key_file, "r") as f:
+            key = f.read()
+        os.environ["WANDB_API_KEY"] = key.strip()
+
+
+def get_wandb_api_key(api_key_file: str = "../config/wandb_api_key.txt") -> str:
+    import os
+
+    try:
+        return os.environ["WANDB_API_KEY"]
+    except KeyError:
+        with open(api_key_file, "r") as f:
+            key = f.read()
+        return key.strip()
+
+
+# Define the function to compute the spectrum
+def spectrum_2d(signal: Array, n_observations: int, normalize: bool = True) -> Array:
+    """This function computes the spectrum of a 2D signal using the Fast Fourier Transform (FFT).
+
+    Parameters
+    ----------
+    signal : an array of shape (T * n_observations * n_observations)
+        A 2D discretized signal represented as a 1D array with shape
+        (T * n_observations * n_observations), where T is the number of time
+        steps and n_observations is the spatial size of the signal.
+
+        T can be any number of channels that we reshape into and
+        n_observations * n_observations is the spatial resolution.
+    n_observations: an integer
+        Number of discretized points. Basically the resolution of the signal.
+    normalize: bool
+        whether to apply normalization to the output of the 2D FFT.
+        If True, normalizes the outputs by ``1/n_observations``
+        (actually ``1/sqrt(n_observations * n_observations)``).
+    Returns
+    --------
+    spectrum: a array
+        A 1D array of shape (s,) representing the computed spectrum.
+        The spectrum is computed using a square approximation to radial
+        binning, meaning that the wavenumber 'bin' into which a particular
+        coefficient is the coefficient's location along the diagonal, indexed
+        from the top-left corner of the 2d FFT output.
+    """
+    T = signal.shape[0]
+    signal = jnp.reshape(signal, (T, n_observations, n_observations))
+
+    if normalize:
+        signal = jnp.fft.fft2(signal, norm="ortho")
+    else:
+        signal = jnp.fft.rfft2(
+            signal, s=(n_observations, n_observations), norm="backward"
+        )
+
+    # 2d wavenumbers following numpy fft convention
+    k_max = n_observations // 2
+    wavenumers = jnp.repeat(
+        jnp.concat(
+            (
+                jnp.arange(start=0, stop=k_max, step=1),
+                jnp.arange(start=-k_max, stop=0, step=1),
+            ),
+            axis=0,
+        ),
+        n_observations,
+        1,
+    )
+    k_x = jnp.transpose(wavenumers, (0, 1))
+    k_y = wavenumers
+
+    # Sum wavenumbers
+    sum_k = jnp.sqrt(k_x**2 + k_y**2)
+    sum_k = sum_k
+
+    # Remove symmetric components from wavenumbers
+    index = -1.0 * jnp.ones((n_observations, n_observations))
+    k_max1 = k_max + 1
+    index[0:k_max1, 0:k_max1] = sum_k[0:k_max1, 0:k_max1]
+
+    spectrum = jnp.zeros((T, n_observations))
+    for j in range(1, n_observations + 1):
+        ind = jnp.where(index == j)
+        spectrum[:, j - 1] = jnp.sum(jnp.abs(signal[:, ind[0], ind[1]]) ** 2, axis=1)
+
+    spectrum = jnp.mean(spectrum, axis=0)
+    return spectrum
+
+
+Number = Union[float, int]
+
+
+def validate_scaling_factor(
+    scaling_factor: Union[None, Number, list[Number], list[list[Number]]],
+    n_dim: int,
+    n_layers: Optional[int] = None,
+) -> Union[None, list[float], list[list[float]]]:
+    """
+    Parameters
+    ----------
+    scaling_factor : None OR float OR list[float] Or list[list[float]]
+    n_dim : int
+    n_layers : int or None; defaults to None
+        If None, return a single list (rather than a list of lists)
+        with `factor` repeated `dim` times.
+    """
+    if scaling_factor is None:
+        return None
+    if isinstance(scaling_factor, (float, int)):
+        if n_layers is None:
+            return [float(scaling_factor)] * n_dim
+
+        return [[float(scaling_factor)] * n_dim] * n_layers
+
+    if (
+        isinstance(scaling_factor, list)
+        and len(scaling_factor) > 0
+        and all([isinstance(s, (float, int)) for s in scaling_factor])
+    ):
+        if n_layers is None and len(scaling_factor) == n_dim:
+            # this is a dim-wise scaling
+            return [float(s) for s in scaling_factor]
+        return [[float(s)] * n_dim for s in scaling_factor]
+
+    if (
+        isinstance(scaling_factor, list)
+        and len(scaling_factor) > 0
+        and all([isinstance(s, (list)) for s in scaling_factor])
+    ):
+        s_sub_pass = True
+        for s in scaling_factor:
+            if all([isinstance(s_sub, (int, float)) for s_sub in s]):
+                pass
+            else:
+                s_sub_pass = False
+            if s_sub_pass:
+                return scaling_factor
+
+    return None
+
+
+def compute_rank(arr: Array) -> Array:
+    # Compute the matrix rank of a array
+    rank = jnp.linalg.matrix_rank(arr)
+    return rank
+
+
+def compute_stable_rank(arr: Array) -> Array:
+    # Compute the stable rank of a array
+    fro_norm = jnp.linalg.norm(arr, ord="fro") ** 2
+    l2_norm = jnp.linalg.norm(arr, ord=2) ** 2
+    rank = fro_norm / l2_norm
+    rank = rank
+    return rank
+
+
+def compute_explained_variance(frequency_max: int, s: Array) -> Array:
+    # Compute the explained variance based on frequency_max and singular
+    # values (s)
+    s_current = s.copy()
+    s_current[frequency_max:] = 0
+    return 1 - jnp.var(s - s_current) / jnp.var(s)
+
+
+def get_project_root() -> Path:
+    root = Path(__file__).parent.parent
+    return root

neojax_operators-0.0.1/pyproject.toml

@@ -0,0 +1,39 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "neojax-operators"
+version = "0.0.1"
+description = "Neural Operators in JAX"
+readme = "README.md"
+requires-python = ">=3.9"
+license = { text = "MIT" }
+authors = [
+  { name = "Paul Gekeler", email = "pgekeler@gmail.com" },
+]
+keywords = [
+  "jax",
+  "neural-operators",
+  "fno",
+  "deep-learning",
+  "scientific-computing",
+]
+classifiers = [
+  "Development Status :: 3 - Alpha",
+  "Intended Audience :: Science/Research",
+  "License :: OSI Approved :: MIT License",
+  "Programming Language :: Python :: 3",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+  "jax",
+  "jaxlib",
+  "equinox",
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["neojax"]
+
+[project.urls]
+Homepage = "https://github.com/paulgekeler/neojax"