Skip to main content
Pre-release

This release is a pre-release and may not be stable for production use.

JaxUtils

CircleCI

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

Metadata

Release files for jaxutils-nightly 0.0.8.dev20230411

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for jaxutils-nightly 0.0.8.dev20230411
File Size Uploaded
jaxutils-nightly-0.0.8.dev20230411.tar.gz 30.1 kB Details

Built distribution (wheel)

Table of built distributions (wheels) for jaxutils-nightly 0.0.8.dev20230411
File Interpreter ABI Platform
jaxutils_nightly-0.0.8.dev20230411-py3-none-any.whl Python 3 none any Details

Total release size: 48.0 kB

Release files / jaxutils-nightly-0.0.8.dev20230411.tar.gz

Download URL jaxutils-nightly-0.0.8.dev20230411.tar.gz
Size 30.1 kB
Tags Source
SHA-256 checksum
How to use checksums
e5a9c8dae0237d02f57994058e37b1b6eab9f11bcd3975d505ac1a3648ce1282
BLAKE2b-256 checksum
How to use checksums
8fa34210ea5d0f4412f0b1bcd3acef40fb4d9a91928c27860352761ea7ebd5ee
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
No
Uploaded via twine/4.0.2 CPython/3.9.0

Release files / jaxutils_nightly-0.0.8.dev20230411-py3-none-any.whl

Download URL jaxutils_nightly-0.0.8.dev20230411-py3-none-any.whl
Size 17.9 kB
Tags Python 3
SHA-256 checksum
How to use checksums
14db708de43ddf4afbd60d7e3510bc1375030719d7c8fce9bfde2d524ef5ccd5
BLAKE2b-256 checksum
How to use checksums
c2f77d910fd47ca9d058f8d9d70fa5ade40b81a9923d3106cc4712c70fe7f4fe
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
No
Uploaded via twine/4.0.2 CPython/3.9.0

Release history Release notifications | RSS feed

This release
Anthropic, PBC Visionary sponsor Bloomberg Visionary sponsor Hudson River Trading Visionary sponsor Meta Visionary sponsor NVIDIA Visionary sponsor Microsoft Sustainability sponsor Depot Continuous Integration AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page