<commit_before># Chex

![CI status](https://github.com/deepmind/chex/workflows/ci/badge.svg)
![docs](https://readthedocs.org/projects/chex/badge/?version=latest)
![pypi](https://img.shields.io/pypi/v/chex)

Chex is a library of utilities for helping to write reliable JAX code.

This includes utils to help:

* Instrument your code (e.g. assertions)
* Debug (e.g. transforming `pmaps` in `vmaps` within a context manager).
* Test JAX code across many `variants` (e.g. jitted vs non-jitted).

## Installation

Chex can be installed with pip directly from github, with the following command:

`pip install git+git://github.com/deepmind/chex.git`

or from PyPI:

`pip install chex`

## Modules Overview

### Dataclass ([dataclass.py](https://github.com/deepmind/chex/blob/master/chex/_src/dataclass.py))

Dataclasses are a popular construct introduced by Python 3.7 to allow to
easily specify typed data structures with minimal boilerplate code. They are
not, however, compatible with JAX and
[dm-tree](https://github.com/deepmind/tree) out of the box.<commit_msg>Recommend installing the PyPI released package first, instead of installing from GitHub, as it may not be as stable as the released version.<commit_after># Chex

![CI status](https://github.com/deepmind/chex/workflows/ci/badge.svg)
![docs](https://readthedocs.org/projects/chex/badge/?version=latest)
![pypi](https://img.shields.io/pypi/v/chex)

Chex is a library of utilities for helping to write reliable JAX code.

This includes utils to help:

* Instrument your code (e.g. assertions)
* Debug (e.g. transforming `pmaps` in `vmaps` within a context manager).
* Test JAX code across many `variants` (e.g. jitted vs non-jitted).

## Installation

You can install the latest released version of Chex from PyPI via:

```sh
pip install chex
```

or you can install the latest development version from GitHub:

```sh
pip install git+https://github.com/deepmind/chex.git
```

## Modules Overview

### Dataclass ([dataclass.py](https://github.com/deepmind/chex/blob/master/chex/_src/dataclass.py))

Dataclasses are a popular construct introduced by Python 3.7 to allow to
easily specify typed data structures with minimal boilerplate code. They are
not, however, compatible with JAX and
[dm-tree](https://github.com/deepmind/tree) out of the box.
{'fraction_non_alphanumeric': 0.11728931364031277, 'fraction_numerical': 0.0017376194613379669, 'mean_word_length': 4.644607843137255, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 11, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '551612', 'n_tokens_mistral': 771, 'n_tokens_neox': 726, 'n_words': 249}