Package: jax-chacha-prng
JAX-ChaCha-PRNGA cryptographically-secure pseudo-random number generator for JAX based on the 20 round ChaCha cipher.The ChaCha cipher was introduced inDaniel J. Bernstein "ChaCha, a variant of Salsa20".The implementation follows the specification in the IRTFRFC 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.APIThe package is split into two modules:Theciphermodule is a full implementation of the ChaCha20 cipher.Therandommodule provides aJAX-style API for the CSPRNG based on the cipher.RandomThe package currently exposes basic RNG functions using the same interface asJAX:chacha.random.PRNGKey: Equivalent tojax.random.PRNGKey: Given a seed of up to 256 bits, it returns aPRNGKeyobject from which randomness can be generated.chacha.random.split: Equivalent tojax.random.split: Splits a givenPRNGKeyinto the desired number of freshPRNGKeyinstances.chacha.random.fold_in: !Deprecated! Equivalent tojax.random.fold_in: Deterministically derives a newPRNGKeyfrom a given one and additional data.chacha.random.random_bits: Equivalent tojax._src.random._random_bits: Raw access to random bits, returned as an array of unsinged integers.chacha.random.uniform: Equivalent tojax.random.uniform: Uniformly sampled floating point numbers in the range[0, 1).Note:PRNGKeyinstances of this ChaCha20-based RNG are not interoperable with those ofjax.random, i.e., you cannot mix them.Security noticeVersions prior to 2.0.0 may repeat random states via thesplitandfold_infunctions.Usage notesPer conventions of pseudo-random number generation in theJAXframework, the functionsrandom_bitsanduniformare
deterministic given the randomness state (thePRNGKeyobject). The user needs to split the state usingsplitbefore each
call torandom_bitsoruniformto get proper pseudo-random numbers. For more details, see what theJAX documentationhas to say about this.Implementation remarksThePRNGKeyobject is the state of the ChaCha cipher with the randomness seed provided by the user as the cipher key.
Every invocation ofsplitgenerates 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 torandom_bitsanduniformwhich can thus output up to 256 GiB of random values
for each state.CipherThe 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.InstallingFor the latest stable version install via pippip install jax-chacha-prngBinaries for glibc based 64-bit linux systems (manylinux wheels) are compiled with CPU and CUDA support (you will have toinstall JAX with CUDA supportto 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 viapip install --no-binary :all: jax-chacha-prngor by installing it directly from thev1-stablebranch:pip install git+https://github.com/DPBayes/jax-chacha-prng@v1-stable#egg=jax-chacha-prngThis 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 ofchacha.native.cuda_supported().Note about JAX versionsJAX 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 thecompatible-jaxinstallation 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.3.25 .
We regularly check the compatible version range, but do not expect new versions of JAX to be immediately tested.VersioningVersion numbers adhere toSemantic Versioning. Changes between releases are tracked inChangeLog.txt.LicenseThe software is licensed under theApache License 2.0.
You can find the full license text inLICENSES/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 theLICENSESdirectory. The repository is fullyREUSE-compliant.AcknowledgementsWe thank the NVIDIA AI Technology Center Finland for their contribution of GPU performance benchmarking and helpful discussions on optimisation.Developing and TestingWe 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 viapytest 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