semvec/pypi-packages huggingface dataset:
jax-toolkit 0.2.0
A collection of jax functions to help with common machine/deep learning related functionality.
License: Apache-2.0
Requires Python: >=3.7

# jax_toolkit

A collection of jax functions to help with common machine/deep learning related functionality. 

[Documentation](https://asmith26.github.io/jax_toolkit/), [PyPi](https://pypi.org/project/jax-toolkit/)

This library currently contains the basics for a number of losses and metrics. We intend to add more complexity and functionality as and when it's needed - of course contributions/pull requests/bug reports etc. are very welcome if you discover problems or need something that is currently missing.

## Installation

```bash
pip install jax_toolkit
```

Or for additional loss function [utils](https://asmith26.github.io/jax_toolkit/losses_and_metrics/#utils):

```bash
pip install jax_toolkit[losses_utils]
```