GitLocker: The Coding Marketplace

Description:

jaxchachaprng 1.4.3

JAX-ChaCha-PRNG
A cryptographically-secure pseudo-random number generator for JAX based on the 20 round ChaCha cipher.
The ChaCha cipher was introduced in Daniel J. Bernstein "ChaCha, a variant of Salsa20".
The implementation follows the specification in the IRTF RFC 7539: "ChaCha20 and Poly1305 for IETF Protocols".
Note that the implementation is not security-hardened. Our threat models assumes
that the machine on which the code is executed is a trusted environment and we
keep key values, cipher states and plaintexts in plain memory.
API
The package is split into two modules:

The cipher module is a full implementation of the ChaCha20 cipher.
The random module provides a JAX-style API for the CSPRNG based on the cipher.

Random
The package currently exposes basic RNG functions using the same interface as JAX:

chacha.random.PRNGKey: Equivalent to jax.random.PRNGKey: Given a seed of up to 256 bits, it returns a PRNGKey object from which randomness can be generated.
chacha.random.split: Equivalent to jax.random.split: Splits a given PRNGKey into the desired number of fresh PRNGKey instances.
chacha.random.fold_in: !Deprecated! Equivalent to jax.random.fold_in: Deterministically derives a new PRNGKey from a given one and additional data.
chacha.random.random_bits: Equivalent to jax._src.random._random_bits: Raw access to random bits, returned as an array of unsinged integers.
chacha.random.uniform: Equivalent to jax.random.uniform: Uniformly sampled floating point numbers in the range [0, 1).

Note: PRNGKey instances of this ChaCha20-based RNG are not interoperable with those of jax.random, i.e., you cannot mix them.
Security notice Versions prior to 2.0.0 may repeat random states via the split and fold_in functions.
Usage notes
Per conventions of pseudo-random number generation in the JAX framework, the functions random_bits and uniform are
deterministic given the randomness state (the PRNGKey object). The user needs to split the state using split before each
call to random_bits or uniform to get proper pseudo-random numbers. For more details, see what the JAX documentation has to say about this.
Implementation remarks
The PRNGKey object is the state of the ChaCha cipher with the randomness seed provided by the user as the cipher key.
Every invocation of split generates random numbers from the given state to use as fresh nonces in the derived states, with the key kept constant.
The counter in the state is used internally within calls to random_bits and uniform which can thus output up to 256 GiB of random values
for each state.
Cipher
The following methods for direct use of the ChaCha20 cipher for encryption are available:
Quick use functions:

chacha.cipher.encrypt_with_key: Encrypt a of any length message by providing a 256 bit key, 96 bit nonce/IV and an optional 32 bit initial counter value.
chacha.cipher.decrypt_with_key: Decrypt a of any length message by providing a 256 bit key, 96 bit nonce/IV and an optional 32 bit initial counter value.

State construction and use:

chacha.cipher.setup_state: Create a ChaCha state structure by providing a 256 bit key, 96 bit nonce/IV and a 32 bit initial counter value.
chacha.cipher.encrypt: Encrypt a message of any length using a ChaCha state structure.
chacha.cipher.decrypt: Decrypt a message of any length using a ChaCha state structure.

Installing
For the latest stable version install via pip
pip install jax-chacha-prng

Binaries for glibc based 64-bit linux systems (manylinux wheels) are compiled with CPU and CUDA support (you will have to install JAX with CUDA support to benefit from this).
Binaries for all other systems are compiled for CPU execution only. This is because JAX does not have CUDA libraries for these systems either.
However, you can instruct pip to instead compile the package from sources via
pip install --no-binary :all: jax-chacha-prng

or by installing it directly from the v1-stable branch:
pip install git+https://github.com/DPBayes/jax-chacha-prng@v1-stable#egg=jax-chacha-prng

This will compile CUDA kernels if the CUDA library is present on the system,
otherwise only CPU kernels will be built. To check whether CUDA kernels were
built and installed, you can check the return value of chacha.native.cuda_supported().
Note about JAX versions
JAX is still under ungoing development and its developers currently give no
guarantee that the API remains stable between releases. However, recent releases
were mostly stable in the interfaces required for JAX-ChaCha-PRNG. In order to allow
usage with JAX-ChaCha-PRNG with the most current JAX release, we therefore do not
currently constrain the JAX version from above in our dependency list.
However, if you should encounter issues with a new JAX release at some point,
you can use the compatible-jax installation target to force usage of the latest
JAX version known to be compatible with JAX-ChaCha-PRNG:
pip install .[compatible-jax]

JAX-ChaCha-PRNG is currently known to work reliably with JAX versions 0.2.12 - 0.4.25 .
We regularly check the compatible version range, but do not expect new versions of JAX to be immediately tested.
Versioning
Version numbers adhere to Semantic Versioning. Changes between releases are tracked in ChangeLog.txt.
License
The software is licensed under the Apache License 2.0.
You can find the full license text in LICENSES/Apache-2.0.txt.
Single files included from third parties may be under a different license, which is annotated in the file
itself and a full license text included in the LICENSES directory. The repository is fully REUSE-compliant.
Acknowledgements
We thank the NVIDIA AI Technology Center Finland for their contribution of GPU performance benchmarking and helpful discussions on optimisation.
Developing and Testing
We welcome any fixes, improvements or other contributions via pull request to this repository.
Before submitting your changes, please make sure to run our Python unit tests via pytest tests/ and
ensure that they all succeed. If you add new functionality, please also add tests.
If you made changes to the native C++/CUDA code, please also compile and run the native tests:
mkdir build
cmake -DBUILD_TESTING=On ..
make -j
./cpu_kernel_tests
./gpu_kernel_tests # if you have CUDA installed and a GPU available