Pre-release
This release is a pre-release and may not be stable for production use.
This project has now been incorporated into GPJax.
JaxUtils
JaxUtils provides utility functions for the JaxGaussianProcesses ecosystem.
Contents
PyTree
Overview
jaxutils.PyTree is a mixin class for registering a python class as a JAX PyTree. You would define your Python class as follows.
class MyClass(jaxutils.PyTree):
...
Example
import jaxutils
from jaxtyping import Float, Array
class Line(jaxutils.PyTree):
def __init__(self, gradient: Float[Array, "1"], intercept: Float[Array, "1"]) -> None
self.gradient = gradient
self.intercept = intercept
def y(self, x: Float[Array, "N"]) -> Float[Array, "N"]
return x * self.gradient + self.intercept
Dataset
Overview
jaxutils.Dataset is a datset abstraction. In future, we wish to extend this to a heterotopic and isotopic data abstraction.
Example
import jaxutils
import jax.numpy as jnp
# Inputs
X = jnp.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
# Outputs
y = jnp.array([[7.0], [8.0], [9.0]])
# Datset
D = jaxutils.Dataset(X=X, y=y)
print(f'The number of datapoints is {D.n}')
print(f'The input dimension is {D.in_dim}')
print(f'The output dimension is {D.out_dim}')
print(f'The input data is {D.X}')
print(f'The output data is {D.y}')
print(f'The data is supervised {D.is_supervised()}')
print(f'The data is unsupervised {D.is_unsupervised()}')
The number of datapoints is 3
The input dimension is 2
The output dimension is 1
The input data is [[1. 2.]
[3. 4.]
[5. 6.]]
The output data is [[7.]
[8.]
[9.]]
The data is supervised True
The data is unsupervised False
You can also add dataset together to concatenate them.
# New inputs
X_new = jnp.array([[1.5, 2.5], [3.5, 4.5], [5.5, 6.5]])
# New outputs
y_new = jnp.array([[7.0], [8.0], [9.0]])
# New dataset
D_new = jaxutils.Dataset(X=X_new, y=y_new)
# Concatenate the two datasets
D = D + D_new
print(f'The number of datapoints is {D.n}')
print(f'The input dimension is {D.in_dim}')
print(f'The output dimension is {D.out_dim}')
print(f'The input data is {D.X}')
print(f'The output data is {D.y}')
print(f'The data is supervised {D.is_supervised()}')
print(f'The data is unsupervised {D.is_unsupervised()}')
The number of datapoints is 6
The input dimension is 2
The output dimension is 1
The input data is [[1. 2. ]
[3. 4. ]
[5. 6. ]
[1.5 2.5]
[3.5 4.5]
[5.5 6.5]]
The output data is [[7.]
[8.]
[9.]
[7.]
[8.]
[9.]]
The data is supervised True
The data is unsupervised False
1 maintainer
Metadata
1 maintainer
Release files for jaxutils-nightly 0.0.8.dev20250316
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| jaxutils-nightly-0.0.8.dev20250316.tar.gz | 30.2 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| jaxutils_nightly-0.0.8.dev20250316-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 48.2 kB
Release files / jaxutils-nightly-0.0.8.dev20250316.tar.gz
| Download URL | jaxutils-nightly-0.0.8.dev20250316.tar.gz |
|---|---|
| Size | 30.2 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
37f720f3cfcc8af7a546e4c7fa88e3fe7d2368315d865927d24a1bff12baac18
|
|
BLAKE2b-256 checksum How to use checksums |
711b35e9dfcb927aba4a5bd5a04f89700930493e7f2c8db60470b7377d2a8ef6
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/6.1.0 CPython/3.9.0
|
Release files / jaxutils_nightly-0.0.8.dev20250316-py3-none-any.whl
| Download URL | jaxutils_nightly-0.0.8.dev20250316-py3-none-any.whl |
|---|---|
| Size | 18.0 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
7a4a6cc93996e3f85005764cb6dd629640598bb5618cd78a23b09cfb41ed7d6f
|
|
BLAKE2b-256 checksum How to use checksums |
51d6b347f73ae4de8f9965acb0f033407d5c543e475450ed44c806920c04f9e2
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/6.1.0 CPython/3.9.0
|
