Repository: DPBayes/jax-chacha-prng (language: python)

Issue:
Support for ARM based systems
Due to its reliance on Intel SIMD intrinsics, the package can currently not be used on any ARM based systems like Apple's M1.

Support for such systems is desired. Therefore the CPU kernels need to be implemented also using ARM NEON intrinsics.


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/build_wheels.yml b/.github/workflows/build_wheels.yml
index 102c824..573fe09 100644
--- a/.github/workflows/build_wheels.yml
+++ b/.github/workflows/build_wheels.yml
@@ -38,7 +38,7 @@ jobs:
         sudo apt-get update
         sudo apt-get install gcc-${{ matrix.cuda-setup[2] }} g++-${{ matrix.cuda-setup[2] }}
         sudo ln -sf /usr/bin/gcc-${{ matrix.cuda-setup[2] }} /usr/bin/gcc
-        sudo ln -sf /usr/bin/g++-${{ matrix.cuda-setup[2] }} /usr/bin/gcc
+        sudo ln -sf /usr/bin/g++-${{ matrix.cuda-setup[2] }} /usr/bin/g++
     - name: Set up CUDA-Toolkit ${{ matrix.cuda-setup[0] }}
       run: |
         wget -q -O installer.run ${{ matrix.cuda-setup[1] }}
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 4b0fa09..07bb7ab 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -1,13 +1,77 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: © 2021 Aalto University
+# SPDX-FileCopyrightText: © 2022 Aalto University
 
 cmake_minimum_required(VERSION 3.18)
 project(jax-chacha20-prng LANGUAGES CXX)
 
 option(BUILD_TESTING "Build tests for native kernels" OFF)
+option(FORCE_GENERIC "Build without CPU architecture optimized instructions" OFF)
 
 set(CMAKE_CXX_STANDARD 14)
-add_compile_options($<$<COMPILE_LANGUAGE:CXX>:-msse>)
+set(cpu_arch_path "${CMAKE_CURRENT_LIST_DIR}/lib/generic/")
+set(cpu_arch_def "ARCH_GENERIC")
+set(SSE_ENABLED "No")
+set(NEON_ENABLED "No")
+
+if(NOT FORCE_GENERIC)
+    try_run(SSE_RUN_SUCCESS SSE_COMPILE_SUCCESS
+        ${CMAKE_CURRENT_BINARY_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/cmake_config/sse_test.cpp
+        #WORKING_DIRECTORY ${CMAKE_CURRENT_LIST_DIR}/cmake_config/
+    )
+
+    if (NOT SSE_COMPILE_SUCCESS)
+        # could not compile sse intrinsics, maybe we need the compiler flag?
+        try_run(SSE_RUN_SUCCESS SSE_COMPILE_SUCCESS
+            ${CMAKE_CURRENT_BINARY_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/cmake_config/sse_test.cpp
+            COMPILE_DEFINITIONS "-msse"
+        )
+        if (SSE_COMPILE_SUCCESS)
+            add_compile_options($<$<COMPILE_LANGUAGE:CXX>:-msse>)
+        endif()
+    endif()
+
+    if (NOT SSE_COMPILE_SUCCESS)
+        try_run(NEON_RUN_SUCCESS NEON_COMPILE_SUCCESS
+            ${CMAKE_CURRENT_BINARY_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/cmake_config/neon_test.cpp
+        )
+
+        if (NOT NEON_COMPILE_SUCCESS)
+            # could not compile neon intrinsics, maybe we need the compiler flag?
+            try_run(NEON_RUN_SUCCESS NEON_COMPILE_SUCCESS
+                ${CMAKE_CURRENT_BINARY_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/cmake_config/neon_test.cpp
+                COMPILE_DEFINITIONS "-mfpu=neon"
+            )
+            if (NEON_COMPILE_SUCCESS)
+                add_compile_options($<$<COMPILE_LANGUAGE:CXX>:-mfpu=neon>)
+            endif()
+        endif()
+    endif()
+
+    if (SSE_COMPILE_SUCCESS)
+        set(cpu_arch_path "${CMAKE_CURRENT_LIST_DIR}/lib/intel/")
+        set(cpu_arch_def "ARCH_INTEL")
+        set(SSE_ENABLED "Yes")
+        if (NOT SSE_RUN_SUCCESS EQUAL 0)
+            message(WARNING "Can compile with Intel SSE instructions but failed a test run (cross-compiling?); will compile with SSE anyways - set FORCE_GENERIC to disable")
+        endif()
+    endif()
+
+    if (NEON_COMPILE_SUCCESS)
+        set(cpu_arch_path "${CMAKE_CURRENT_LIST_DIR}/lib/arm/")
+        set(cpu_arch_def "ARCH_ARM")
+        set(NEON_ENABLED "Yes")
+        if (NOT NEON_RUN_SUCCESS EQUAL 0)
+            message(WARNING "Can compile with ARM Neon instructions but failed a test run (cross-compiling?); will compile with Neon anyways - set FORCE_GENERIC to disable")
+        endif()
+    endif()
+
+endif()
+
+
+message("-- Detected architecture - " ${CMAKE_SYSTEM_PROCESSOR})
+
+message("-- SSE instructions enabled - ${SSE_ENABLED}")
+message("-- ARM Neon instructions enabled - ${NEON_ENABLED}")
 
 # native kernels require pybind11; however, if it is installed from pip, it has issues
 # in detecting the Python interpreter and libraries during CMake build in a conda env,
@@ -34,11 +98,16 @@ if(NOT EXISTS "${PROJECT_SOURCE_DIR}/extern/pybind11/CMakeLists.txt")
     message(FATAL_ERROR "The pybind11 repository was not downloaded! Please run manually: git submodule update --init --recursive .")
 endif()
 
-find_package(Python COMPONENTS Interpreter Development REQUIRED)
+find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
 add_subdirectory(extern/pybind11)
 # find_package(pybind11 REQUIRED) # broken; does not use correct Python interpreter/libraries in a conda env
 
-find_package(OpenMP REQUIRED)
+find_package(OpenMP)
+if(OpenMP_CXX_FOUND)
+    add_compile_definitions(OPENMP_AVAILABLE)
+else()
+    message(WARNING "OpenMP not found - Compiling without, but you may see lower performance.")
+endif()
 
 include_directories(${CMAKE_CURRENT_LIST_DIR}/lib)
 
@@ -46,15 +115,24 @@ pybind11_add_module(
     native
     ${CMAKE_CURRENT_LIST_DIR}/lib/cpu_kernel.cpp
     ${CMAKE_CURRENT_LIST_DIR}/lib/python_bindings.cpp)
-target_link_libraries(native PUBLIC OpenMP::OpenMP_CXX)
+target_include_directories(native PRIVATE ${cpu_arch_path})
+target_compile_definitions(native PRIVATE ${cpu_arch_def})
+
+if(OpenMP_CXX_FOUND)
+    target_link_libraries(native PUBLIC OpenMP::OpenMP_CXX)
+endif()
 
 if (BUILD_TESTING)
     add_executable(cpu_kernel_tests
         ${CMAKE_CURRENT_LIST_DIR}/tests/lib/cpu_kernel_tests.cpp
         ${CMAKE_CURRENT_LIST_DIR}/lib/cpu_kernel.cpp)
     target_include_directories(cpu_kernel_tests PRIVATE
-        ${CMAKE_CURRENT_LIST_DIR}/lib/)
-    target_link_libraries(cpu_kernel_tests PUBLIC OpenMP::OpenMP_CXX)
+        ${CMAKE_CURRENT_LIST_DIR}/lib/
+        ${cpu_arch_path})
+    target_compile_definitions(cpu_kernel_tests PRIVATE ${cpu_arch_def})
+    if (OpenMP_CXX_FOUND)
+        target_link_libraries(cpu_kernel_tests PUBLIC OpenMP::OpenMP_CXX)
+    endif()
 endif()
 
 include(CheckLanguage)
@@ -78,7 +156,11 @@ if (CMAKE_CUDA_COMPILER)
         target_include_directories(gpu_kernel_tests PRIVATE
             ${CMAKE_CURRENT_LIST_DIR}/lib/)
         target_link_libraries(gpu_kernel_tests PUBLIC OpenMP::OpenMP_CXX)
+        target_compile_definitions(gpu_kernel_tests PRIVATE ${cpu_arch_def})
         set_property(TARGET gpu_kernel_tests PROPERTY CUDA_ARCHITECTURES 35)
+        if (OpenMP_CXX_FOUND)
+            target_link_libraries(gpu_kernel_tests PUBLIC OpenMP::OpenMP_CXX)
+        endif()
     endif()
 else()
     message(WARNING "CUDA not found - building for CPU only!")
diff --git a/ChangeLog.txt b/ChangeLog.txt
index ce3787c..df20bea 100644
--- a/ChangeLog.txt
+++ b/ChangeLog.txt
@@ -1,19 +1,32 @@
+- 2.0.0-rc.2:
+    - Fix: OpenMP is no longer a strict requirement for installation.
+    - Added: chacha.native.openmp_accelerated, returns True if CPU kernels are parallelised using OpenMP.
+    - Added: Generic (not hardware-specific) CPU kernels as well as vectorized CPU kernels for ARM CPUs.
+    - Added: random.is_state_invalidated to check if a PRNGKey was invalidated.
+    - Changed: random.split now can perform nested splits only up to a limit and returns
+        an invalidated state if this limit is exceeded. This is to guarantee
+        unique states after each split derived from the same inital randomness state.
+    - Changed: random.random_bits now return a new rng key for subsequent calls to
+        produce randmoness as well as an error flag integer,
+        which is set if the input randomness state was invalidated or the
+        randomness counter overflows.
+    - Changed: random.uniform now returns all NaNs if the input randomness state
+        was invalidated or the randomness counter overflows. It also accepts additional
+        kwarg return_next_key that returns a new rng key.
+    - Removed: random.fold_in.
+- 1.2.0:
+    - Deprecated: random.fold_in.
 - 1.1.1:
     - Added: support for jax up to v0.3.13.
 - 1.1.0:
     - Added: support for multiple layers of vmap wraps for functions in chacha.random.
     - Added: support for jax up to v0.3.1.
 - 1.0.0:
+    - Added: Native implementations for ChaCha block function as JAX ops.
     - Added: chacha.native.cuda_supported, returns True if cuda kernels were compiled.
-    - Added: support for jax up to v0.2.27.
-- 1.0.0-rc.3:
     - Added: wheels for Python 3.10.
-- 1.0.0-rc.2:
-    - Added: support for jax v0.2.21 - v0.2.22.
+    - Added: support for jax up to v0.2.27.
     - Removed: support for jax v0.2.10 - v0.2.11.
-- 1.0.0-rc.1:
-    - Added: support for jax v0.2.15 - v0.2.20.
-    - Added: Native implementations for ChaCha block function as JAX ops.
 - 0.1.0-alpha.2:
     - Fix: random.random_bits() no longer fails for empty shape.
     - Added: support for jax v0.2.14.
diff --git a/README.md b/README.md
index dc57516..03f959f 100644
--- a/README.md
+++ b/README.md
@@ -28,6 +28,11 @@ The package currently exposes basic RNG functions using the same interface as `J
 
 *Note*: `PRNGKey` instances of this ChaCha20-based RNG are not interoperable with those of `jax.random`, i.e., you cannot mix them.
 
+`random_bits` and (optionally) `uniform` return a new rng key resulting from advancing the counter field in the key provided to them
+according to the amount of randomness they were tasked to produce. This key may be used for subsequent calls to producing randomness,
+but CANNOT be used as an input for `split`.
+
+
 **Security notice**: Versions prior to 2.0.0 may repeat random states via the `split` and `fold_in` functions and
 therefore not produce sufficiently random outputs.
 
@@ -58,16 +63,23 @@ State construction and use:
 
 ## Installing
 
-For the latest stable version install from the `v1-stable` branch via `pip`:
+For the latest stable version install from the `stable` branch via `pip`:
+```
+pip install git+https://github.com/DPBayes/jax-chacha-prng@stable#egg=jax-chacha-prng
+```
+
+For the latest stable release of major version X, use the `vX-stable` branch, i.e.,
+for major version 2 use
 ```
-pip install git+https://github.com/DPBayes/jax-chacha-prng@v1-stable#egg=jax-chacha-prng
+pip install git+https://github.com/DPBayes/jax-chacha-prng@v2-stable#egg=jax-chacha-prng
 ```
 
 Installation 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()`.
 
-Pre-built binary wheels are also available alongside the releases on GitHub.
+Pre-built binary wheels are also available alongside the releases on GitHub, however,
+these are currently only available for the x86_64 platform.
 
 ### Note about JAX versions
 
@@ -84,7 +96,7 @@ 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.2.27 .
+JAX-ChaCha-PRNG is currently known to work reliably with JAX versions 0.2.12 - 0.3.17 .
 We regularly check the compatible version range, but do not expect new versions of JAX to be immediately tested.
 
 ## Versioning
diff --git a/arm_container/Dockerfile b/arm_container/Dockerfile
new file mode 100644
index 0000000..9b6926b
--- /dev/null
+++ b/arm_container/Dockerfile
@@ -0,0 +1,23 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: © 2022 Aalto University
+
+FROM ubuntu:jammy AS qemu-arm-base
+
+RUN sed -i 's/http/[arch=i386,amd64] http/g' /etc/apt/sources.list
+RUN echo "deb [arch=arm64] http://ports.ubuntu.com/ jammy main restricted" >> /etc/apt/sources.list
+RUN echo "deb [arch=arm64] http://ports.ubuntu.com/ jammy-updates main restricted" >> /etc/apt/sources.list
+
+RUN dpkg --add-architecture arm64
+RUN apt update && apt upgrade -y && apt install -y qemu binfmt-support qemu-user-static libc6:arm64
+
+RUN /bin/bash
+
+FROM qemu-arm-base
+
+RUN apt install -y gcc-aarch64-linux-gnu g++-aarch64-linux-gnu libc6-dev-arm64-cross cmake python3 python3-dev libpython3-dev:arm64
+RUN mkdir -p /work/build/ /work/src/
+ADD docker_main.sh /work/docker_main.sh
+RUN chmod +x /work/docker_main.sh
+WORKDIR /work
+
+ENTRYPOINT [ "/work/docker_main.sh" ]
diff --git a/arm_container/docker_main.sh b/arm_container/docker_main.sh
new file mode 100644
index 0000000..4d60967
--- /dev/null
+++ b/arm_container/docker_main.sh
@@ -0,0 +1,17 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: © 2022 Aalto University
+
+cd /work/build
+cmake \
+    -DCMAKE_SYSTEM_NAME=Linux -DCMAKE_SYSTEM_PROCESSOR=aarch64 \
+    -DCMAKE_C_COMPILER=/usr/bin/aarch64-linux-gnu-gcc \
+    -DCMAKE_CXX_COMPILER=/usr/bin/aarch64-linux-gnu-g++ \
+    -DCMAKE_FIND_ROOT_PATH=/usr/aarch64-linux-gnu/ \
+    -DCMAKE_FIND_ROOT_PATH_MODE_PROGRAM=BOTH \
+    -DCMAKE_CROSSCOMPILING_EMULATOR="/usr/bin/qemu-aarch64-static;-L;/usr/aarch64-linux-gnu/" \
+    -DBUILD_TESTING=On \
+    -DFORCE_GENERIC=$FORCE_GENERIC \
+    -DCMAKE_BUILD_TYPE=Release ../src/
+make
+/usr/bin/qemu-aarch64-static -L /usr/aarch64-linux-gnu/ ./cpu_kernel_tests
\ No newline at end of file
diff --git a/chacha/cipher.py b/chacha/cipher.py
index d59d41c..943f8cc 100644
--- a/chacha/cipher.py
+++ b/chacha/cipher.py
@@ -73,12 +73,12 @@ def setup_state(
     elif jax.lax.dtype(counter) != ChaChaStateElementType or jnp.size(counter) != ChaChaCounterSizeInWords:
         raise ValueError("counter must be a single 32-bit unsigned integer!")
 
-    key_bits = key_array.size * 4
-    if key_bits == 16:
+    key_bytes = key_array.size * 4
+    if key_bytes == 16:
         key_array = jnp.tile(key_array, 2)
     key_array = key_array.reshape(2, 4)  # type: ignore # numpy seems confused about this
     inputs = jnp.hstack((counter, iv_array))
-    state = ChaChaState(jnp.vstack((KEY_GEN_CONSTANTS[key_bits], key_array, inputs)))
+    state = ChaChaState(jnp.vstack((KEY_GEN_CONSTANTS[key_bytes], key_array, inputs)))
     return state
 
 
diff --git a/chacha/native.pyi b/chacha/native.pyi
index a606f12..bd9050b 100644
--- a/chacha/native.pyi
+++ b/chacha/native.pyi
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: © 2021 Aalto University
+# SPDX-FileCopyrightText: © 2022 Aalto University
 
 from typing import Callable, Sequence, Any
 
@@ -9,3 +9,4 @@ XlaCustomCallGPU = Callable[[Any, Sequence[bytes], bytes, int], None]
 def cpu_chacha20_block_factory() -> XlaCustomCallCPU: ...
 def gpu_chacha20_block_factory() -> XlaCustomCallGPU: ...
 def cuda_supported() -> bool: ...
+def openmp_accelerated() -> bool: ...
diff --git a/chacha/random.py b/chacha/random.py
index 6d7e794..b0830d3 100644
--- a/chacha/random.py
+++ b/chacha/random.py
@@ -10,8 +10,8 @@
 
 The following invariants hold:
 - The 256 bit key provides base randomness that is expanded by PRNG; it given by the user as seed to function `PRNGKey`
-- The 32 bit counter in a randomness state is always set to zero; randomness expander such as `random_bits` increment
-    it internally to provide streams of randomness.
+- The 32 bit counter in a randomness state is always incremented by randomness expanders such as `random_bits`
+    to provide streams of randomness.
 - The 96 bit IV is used for randomness state splits using the `split` function; splitting results in a new state
     that maintains the same cipher key, a counter value of zero and a fresh IV derived from the previous IV in such
     a way that the new IV is unique among all states derived from the same initial key state. An IV value of zero
@@ -70,7 +70,7 @@ def is_state_invalidated(rng_key: RNGState) -> bool:
 
 @partial(jax.jit, static_argnums=(1, 2))
 def random_bits(rng_key: RNGState, bit_width: int, shape: typing.Sequence[int])\
-        -> typing.Tuple[jnp.ndarray, jnp.uint32]:
+        -> typing.Tuple[jnp.ndarray, RNGState, jnp.uint32]:
     """ Generates an array containing random integers.
 
     Note that this function enters a failure state if the `rng_key` is invalidated
@@ -79,13 +79,14 @@ def random_bits(rng_key: RNGState, bit_width: int, shape: typing.Sequence[int])\
     of containing an array of random values.
 
     Args:
-      rng_key: The not-invalidated PRNGKey object from which to generate random bits.
+      rng_key: The not-invalidated RNGState object from which to generate random bits.
       bit_width: The number of bits in each element of the output.
       shape: The shape of the output array.
 
     Returns:
       A tuple containing
         - array of the given shape containing uniformly random unsigned integers with the given bit width,
+        - the next RNGState after generating the requested amount of random bits (with the counter value advanced),
         - an integer representing an array of error flags (see `ErrorFlag`)
     """
     if bit_width not in _UINT_DTYPES:
@@ -106,13 +107,17 @@ def generate_block(c: RNGState) -> jnp.ndarray:
     out = cc.serialize(blocks, dtype)
     assert jnp.size(out) >= size
 
-    counter_exceeded = cc.get_counter(rng_key) >= jnp.uint32(-num_blocks)
+    next_rng_key = cc.increase_counter(rng_key, num_blocks)
+
+    counter_exceeded = cc.get_counter(rng_key) >= cc.get_counter(next_rng_key)  # detect wrap-around of counter
     error_flags = jnp.uint32((is_state_invalidated(rng_key) << 1) ^ counter_exceeded)
 
     out = out[:size].reshape(shape)
     out = jnp.where(error_flags == 0, out, out ^ out)
 
-    return out, error_flags
+    next_rng_key = jnp.where(error_flags == 0, next_rng_key, jnp.zeros_like(rng_key))
+
+    return out, next_rng_key, error_flags
 
 
 @partial(jax.jit, static_argnums=(1,))
@@ -130,28 +135,28 @@ def _split(rng_key: RNGState, num: int) -> RNGState:
     assert new_nonce_base.shape == (defs.ChaChaNonceSizeInWords,)
 
     split_nesting_exceeded = (old_nonce[0] >= (1 << (32 - bitlength_num))) | jnp.all(old_nonce == 0)
+    state_previuosly_used = cc.get_counter(rng_key) > 0
 
     def make_rng_key(i: int) -> RNGState:
         nonce = jnp.concatenate((new_nonce_base[:defs.ChaChaNonceSizeInWords - 1], new_nonce_base[-1:] ^ i))
         nonce *= (1 - split_nesting_exceeded)  # set the nonce to 0 (invalid state) if split nesting limit is exceeded
+        nonce *= (1 - state_previuosly_used) # set the nonced to 0 (invalid state) if the state was already used to generate randomness
         return cc.set_counter(cc.set_nonce(rng_key, nonce), 0)
 
     return jax.vmap(make_rng_key)(jnp.arange(num, dtype=rng_key.dtype))
 
-# TODO: deprecate fold_in for v1.x update; release changed splitting as v2
 
-
-@partial(jax.jit, static_argnums=(1, 2))
+@partial(jax.jit, static_argnums=(1, 2, 5))
 def _uniform(
         rng_key: RNGState,
         shape: typing.Tuple[int],
         dtype: type,
         minval: jnp.float_,
-        maxval: jnp.float_
+        maxval: jnp.float_,
+        return_next_key: bool
     ) -> jnp.ndarray:  # noqa:E121,E125
     _check_shape("uniform", shape)
     if not jnp.issubdtype(dtype, np.floating):
-        print("encountered exc in _uniform")
         raise TypeError("uniform only accepts floating point dtypes.")
 
     minval = jax.lax.convert_element_type(minval, dtype)
@@ -164,7 +169,7 @@ def _uniform(
 
     assert nbits in (16, 32, 64)
 
-    bits, errors = random_bits(rng_key, nbits, shape)
+    bits, next_rng_key, errors = random_bits(rng_key, nbits, shape)
 
     # The strategy here is to randomize only the mantissa bits with an exponent of
     # 1 (after applying the bias), then shift and scale to the desired range. The
@@ -180,7 +185,10 @@ def _uniform(
         jax.lax.reshape(floats * (maxval - minval) + minval, shape)
     )
 
-    return jnp.where(errors == 0, result, result * jnp.nan)
+    result = jnp.where(errors == 0, result, result * jnp.nan)
+    if return_next_key:
+        return result, next_rng_key
+    return result
 
 
 def PRNGKey(seed: typing.Union[jnp.ndarray, int, bytes]) -> RNGState:
@@ -241,7 +249,8 @@ def uniform(
         shape: typing.Sequence[int] = (),
         dtype: np.dtype = jnp.float64,
         minval: typing.Union[float, jnp.ndarray] = 0.,
-        maxval: typing.Union[float, jnp.ndarray] = 1.
+        maxval: typing.Union[float, jnp.ndarray] = 1.,
+        return_next_key: bool = False
     ) -> jnp.ndarray:  # noqa:E121,E125
     """Samples uniform random values in [minval, maxval) with given shape/dtype.
 
@@ -249,18 +258,20 @@ def uniform(
     the output will be an array of the requested size, containing only NaN values.
 
     Args:
-      key: The PRNGKey.
+      key: The RNGState.
       shape: An optional tuple of nonnegative integers representing the result shape.
       dtype: An optional float dtype for the returned values (default float64).
-      minval: An ptional minimum (inclusive) value broadcast-compatible with shape for the range (default 0).
+      minval: An optional minimum (inclusive) value broadcast-compatible with shape for the range (default 0).
       maxval: An optional maximum (exclusive) value broadcast-compatible with shape for the range (default 1).
+      return_next_key: An optional boolean flag. If `True`, the function returns a new RNGState (with advanced counter).
 
     Returns:
-      A random array with the specified shape and dtype.
+      A random array with the specified shape and dtype if `return_next_key` is `False`.
+      A tuple consisting of the random array and a new RNGState if `return_next_key` is `True`.
 
     """
     if not jax.dtypes.issubdtype(dtype, np.floating):
         raise TypeError(f"dtype argument to `uniform` must be a float dtype, got {dtype}")
     dtype = jax.dtypes.canonicalize_dtype(dtype)
     shape = _canonicalize_shape(shape)
-    return _uniform(key, shape, dtype, minval, maxval)
+    return _uniform(key, shape, dtype, minval, maxval, return_next_key)
diff --git a/chacha/version.py b/chacha/version.py
index 85ac9a6..d5812e4 100644
--- a/chacha/version.py
+++ b/chacha/version.py
@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: © 2021,2022 Aalto University
+# SPDX-FileCopyrightText: © 2022 Aalto University
 
 MAJOR_VERSION = 2
 MINOR_VERSION = 0
 PATCH_VERSION = 0
-EXT_VERSION = ""
+EXT_VERSION = "rc.2"
 
 EXT_VERSION_SUFFIX = f"-{EXT_VERSION}" if len(EXT_VERSION) > 0 else ""
 
diff --git a/lib/arm/cpu_kernel_arch.hpp b/lib/arm/cpu_kernel_arch.hpp
new file mode 100644
index 0000000..fcf39aa
--- /dev/null
+++ b/lib/arm/cpu_kernel_arch.hpp
@@ -0,0 +1,72 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: © 2022 Aalto University
+
+#pragma once
+
+#include "../defs.hpp"
+#include <stdint.h>
+#include <utility>
+#include <arm_neon.h>
+
+
+template <uint num_positions>
+inline uint32x4_t rotate_vector_left(uint32x4_t vec)
+{
+    return vextq_u32(vec, vec, num_positions);
+}
+
+template <>
+inline uint32x4_t rotate_vector_left<0>(uint32x4_t vec)
+{
+    return vec;
+}
+
+struct StateRow
+{
+private:
+    uint32x4_t values;
+
+public:
+    StateRow() { }
+    StateRow(uint32x4_t vec) : values(std::move(vec)) { }
+    StateRow(const uint32_t row_values[ChaChaStateWordsPerRow]) : values(vld1q_u32(row_values)) { }
+
+    inline StateRow& operator+=(const StateRow other)
+    {
+        values = vaddq_u32(values, other.values);
+        return *this;
+    }
+
+    inline StateRow& operator^=(const StateRow other)
+    {
+        values = veorq_u32(values, other.values);
+        return *this;
+    }
+
+    inline StateRow& operator<<=(int num_bits)
+    {
+        values = vorrq_u32(
+            vshlq_n_u32(values, num_bits),
+            vshrq_n_u32(values, 32 - num_bits)
+        ); // (value << num_bits) ^ (value >> (32 - num_bits));
+        return *this;
+    }
+
+    template <uint num_positions>
+    inline StateRow rotate_elements_left() const
+    {
+        return StateRow(rotate_vector_left<num_positions>(values));
+    }
+
+    template <uint num_positions>
+    inline StateRow rotate_elements_right() const
+    {
+        return rotate_elements_left<(ChaChaStateWordsPerRow - num_positions) % ChaChaStateWordsPerRow>();
+    }
+
+    inline void unvectorize(uint32_t out_buffer[ChaChaStateWordsPerRow]) const
+    {
+        vst1q_u32(out_buffer, values);
+    }
+
+};
diff --git a/lib/cpu_kernel.cpp b/lib/cpu_kernel.cpp
index 8ef44fa..1ca5742 100644
--- a/lib/cpu_kernel.cpp
+++ b/lib/cpu_kernel.cpp
@@ -1,101 +1,106 @@
 // SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: © 2021 Aalto University
+// SPDX-FileCopyrightText: © 2022 Aalto University
 
 #include "cpu_kernel.hpp"
 
-inline __m128i rotate_left(__m128i values, uint num_bits)
+VectorizedState::VectorizedState(StateRow a, StateRow b, StateRow c, StateRow d) : rows{a, b, c, d} { }
+
+VectorizedState::VectorizedState(const uint32_t state[ChaChaStateSizeInWords]) : rows{
+    StateRow(state + 0 * ChaChaStateSizeInRows),
+    StateRow(state + 1 * ChaChaStateSizeInRows),
+    StateRow(state + 2 * ChaChaStateSizeInRows),
+    StateRow(state + 3 * ChaChaStateSizeInRows)
+} { }
+
+void VectorizedState::unvectorize(uint32_t out_state[ChaChaStateSizeInWords]) const
 {
-    return _mm_xor_si128(
-        _mm_slli_epi32(values, num_bits),
-        _mm_srli_epi32(values, 32 - num_bits)
-    ); // (value << num_bits) ^ (value >> (32 - num_bits));
+    for (uint i = 0; i < ChaChaStateSizeInRows; ++i)
+    {
+        rows[i].unvectorize(out_state + i * ChaChaStateSizeInRows);
+    }
 }
 
-VectorizedState quarterround_sse(VectorizedState state)
+StateRow& VectorizedState::operator[](uint i)
 {
-    __m128i a = state[0];
-    __m128i b = state[1];
-    __m128i c = state[2];
-    __m128i d = state[3];
-    a = _mm_add_epi32(a, b); // a += b;
-    d = _mm_xor_si128(d, a); // d ^= a;
-    d = rotate_left(d, 16);
-    c = _mm_add_epi32(c, d); // c += d;
-    b = _mm_xor_si128(b, c); // b ^= c;
-    b = rotate_left(b, 12);
-    a = _mm_add_epi32(a, b); // a += b;
-    d = _mm_xor_si128(d, a); // d ^= a;
-    d = rotate_left(d, 8);
-    c = _mm_add_epi32(c, d); // c += d;
-    b = _mm_xor_si128(b, c); // b ^= c;
-    b = rotate_left(b, 7);
-    return VectorizedState(a, b, c, d);
+    return rows[i];
 }
 
-void pack_diagonals(VectorizedState& out_state, VectorizedState in_state)
+VectorizedState& VectorizedState::operator+=(VectorizedState other)
 {
-    out_state[0] = rotate_elements_left<0>(in_state[0]);
-    out_state[1] = rotate_elements_left<1>(in_state[1]);
-    out_state[2] = rotate_elements_left<2>(in_state[2]);
-    out_state[3] = rotate_elements_left<3>(in_state[3]);
+    for (uint i = 0; i < ChaChaStateSizeInRows; ++i)
+    {
+        rows[i] += other.rows[i];
+    }
+    return *this;
 }
 
-void unpack_diagonals(VectorizedState& out_state, VectorizedState in_state)
+VectorizedState VectorizedState::operator+(VectorizedState other) const
 {
-    out_state[0] = rotate_elements_right<0>(in_state[0]);
-    out_state[1] = rotate_elements_right<1>(in_state[1]);
-    out_state[2] = rotate_elements_right<2>(in_state[2]);
-    out_state[3] = rotate_elements_right<3>(in_state[3]);
+    VectorizedState result(*this);
+    result += other;
+    return result;
 }
 
-VectorizedState double_round_sse(VectorizedState state)
+/// This implements what Bernstein calls a quarterround, but does so in a
+/// vectorized manner, i.e., it performs all quarterrounds over the
+/// state matrix's rows concurrently.
+VectorizedState round(VectorizedState state)
 {
-    state = quarterround_sse(state);
-    pack_diagonals(state, state);
-    state = quarterround_sse(state);
-    unpack_diagonals(state, state);
-
-    return state;
+    StateRow a = state[0];
+    StateRow b = state[1];
+    StateRow c = state[2];
+    StateRow d = state[3];
+    a += b;
+    d ^= a;
+    d <<= 16;
+    c += d;
+    b ^= c;
+    b <<= 12;
+    a += b;
+    d ^= a;
+    d <<= 8;
+    c += d;
+    b ^= c;
+    b <<= 7;
+    return VectorizedState(a, b, c, d);
 }
 
-VectorizedState add_states_sse(VectorizedState x, VectorizedState y)
+void pack_diagonals(VectorizedState& out_state, VectorizedState in_state)
 {
-    VectorizedState out;
-    for (uint i = 0; i < 4; ++i)
-    {
-        out[i] = _mm_add_epi32(x[i], y[i]);
-    }
-    return out;
+    out_state[0] = in_state[0].rotate_elements_left<0>();
+    out_state[1] = in_state[1].rotate_elements_left<1>();
+    out_state[2] = in_state[2].rotate_elements_left<2>();
+    out_state[3] = in_state[3].rotate_elements_left<3>();
 }
 
-VectorizedState vectorize_state(const uint32_t state[16])
+void unpack_diagonals(VectorizedState& out_state, VectorizedState in_state)
 {
-    VectorizedState vec_state;
-    for (uint i = 0; i < 4; ++i)
-    {
-        vec_state[i] = _mm_load_si128(&(reinterpret_cast<const __m128i*>(state)[i]));
-    }
-    return vec_state;
+    out_state[0] = in_state[0].rotate_elements_right<0>();
+    out_state[1] = in_state[1].rotate_elements_right<1>();
+    out_state[2] = in_state[2].rotate_elements_right<2>();
+    out_state[3] = in_state[3].rotate_elements_right<3>();
 }
 
-void unvectorize_state(uint32_t out_state[16], VectorizedState vec_state)
+VectorizedState double_round(VectorizedState state)
 {
-    for (uint i = 0; i < 4; ++i)
-    {
-        _mm_store_si128(&(reinterpret_cast<__m128i*>(out_state)[i]), vec_state[i]);
-    }
+    state = round(state);
+    pack_diagonals(state, state);
+    state = round(state);
+    unpack_diagonals(state, state);
+
+    return state;
 }
 
-void chacha20_block_sse(uint32_t out_state[16], const uint32_t in_state[16])
+void chacha20_block(uint32_t out_state[16], const uint32_t in_state[16])
 {
-    VectorizedState vec_in_state = vectorize_state(in_state);
-    VectorizedState vec_tmp_state = double_round_sse(vec_in_state);
+    VectorizedState vec_in_state(in_state);
+    VectorizedState vec_tmp_state = double_round(vec_in_state);
     for (uint i = 0; i < ChaChaDoubleRoundCount - 1; ++i)
     {
-        vec_tmp_state = double_round_sse(vec_tmp_state);
+        vec_tmp_state = double_round(vec_tmp_state);
     }
-    vec_tmp_state = add_states_sse(vec_in_state, vec_tmp_state);
-    unvectorize_state(out_state, vec_tmp_state);
+    vec_tmp_state += vec_in_state;
+    vec_tmp_state.unvectorize(out_state);
 }
 
 void cpu_chacha20_block(void* out_buffer, const void** in_buffers)
@@ -103,10 +108,12 @@ void cpu_chacha20_block(void* out_buffer, const void** in_buffers)
     uint32_t num_states = *reinterpret_cast<const uint32_t*>(in_buffers[0]);
     const uint32_t* in_states = reinterpret_cast<const uint32_t*>(in_buffers[1]);
     uint32_t* out_state = reinterpret_cast<uint32_t*>(out_buffer);
+    #ifdef OPENMP_AVAILABLE
     #pragma omp parallel for
+    #endif
     for (uint32_t i = 0; i < num_states; ++i)
     {
         uint32_t offset = ChaChaStateSizeInWords * i;
-        chacha20_block_sse(out_state + offset, in_states + offset);
+        chacha20_block(out_state + offset, in_states + offset);
     }
 }
diff --git a/lib/cpu_kernel.hpp b/lib/cpu_kernel.hpp
index b58d2b5..87dd845 100644
--- a/lib/cpu_kernel.hpp
+++ b/lib/cpu_kernel.hpp
@@ -1,70 +1,60 @@
 // SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: © 2021 Aalto University
+// SPDX-FileCopyrightText: © 2022 Aalto University
 
 #pragma once
 
-#include <stdint.h>
-#include <immintrin.h>
-
 #include "defs.hpp"
 
+#include <cpu_kernel_arch.hpp>
 
-struct VectorizedState
-{
-    __m128i values[4];
-    VectorizedState() = default;
-    VectorizedState(__m128i a, __m128i b, __m128i c, __m128i d) : values{a, b, c, d} { }
-    __m128i& operator[](uint i) { return values[i]; }
-};
 
-__m128i rotate_left(__m128i values, uint num_bits);
-
-VectorizedState quarterround_sse(VectorizedState state);
-
-// Wrap _mm_shuffle_epi32 in a template to enforce that rotation_immediate
-// is always known at compile time.
-template <uint8_t rotation_immediate>
-inline __m128i _mm_shuffle_epi32_templated(__m128i val)
+inline StateRow operator+(StateRow left, StateRow right)
 {
-    return _mm_shuffle_epi32(val, rotation_immediate);
+    StateRow result(left);
+    result += right;
+    return result;
 }
 
-// Rotate elements in a 4-vec
-template <uint num_positions>
-inline __m128i rotate_elements_left(__m128i vec)
+inline StateRow operator^(StateRow left, StateRow right)
 {
-    constexpr uint8_t rotation_lookup[4] = {
-        0b11100100,
-        0b00111001,
-        0b01001110,
-        0b10010011
-    };
-    constexpr uint8_t rotation_immediate = rotation_lookup[num_positions];
-    // using the templated wrapper for _mm_shuffle_epi32, otherwise
-    // gcc may be confused and decide to not treat rotation_immediate as
-    // compile-time known in debug mode (-O0) for some reason, resulting
-    // in errors from _mm_shuffle_epi32:
-    return _mm_shuffle_epi32_templated<rotation_immediate>(vec);
+    StateRow result(left);
+    result ^= right;
+    return result;
 }
 
-template <uint num_positions>
-inline __m128i rotate_elements_right(__m128i vec)
+inline StateRow operator<<(StateRow row, uint num_bits)
 {
-    return rotate_elements_left<(4 - num_positions) % 4>(vec);
+    StateRow result(row);
+    result <<= num_bits;
+    return result;
 }
 
-void pack_diagonals(VectorizedState& out_state, VectorizedState in_state);
+template <>
+inline StateRow StateRow::rotate_elements_left<0>() const
+{
+    return *this;
+}
 
-void unpack_diagonals(VectorizedState& out_state, VectorizedState in_state);
 
-VectorizedState double_round_sse(VectorizedState state);
+struct VectorizedState
+{
+private:
+    StateRow rows[ChaChaStateSizeInRows];
 
-VectorizedState add_states_sse(VectorizedState x, VectorizedState y);
+public:
+    VectorizedState() = default;
+    VectorizedState(StateRow a, StateRow b, StateRow c, StateRow d);
+    VectorizedState(const uint32_t state[ChaChaStateSizeInWords]);
 
-VectorizedState vectorize_state(const uint32_t state[16]);
+    void unvectorize(uint32_t out_state[ChaChaStateSizeInWords]) const;
 
-void unvectorize_state(uint32_t out_state[16], VectorizedState vec_state);
+    StateRow& operator[](uint i);
 
-void chacha20_block_sse(uint32_t out_state[16], const uint32_t in_state[16]);
+    VectorizedState& operator+=(VectorizedState other);
+    VectorizedState operator+(VectorizedState other) const;
+};
 
+void unpack_diagonals(VectorizedState& out_state, VectorizedState in_state);
+void pack_diagonals(VectorizedState& out_state, VectorizedState in_state);
+void chacha20_block(uint32_t out_state[16], const uint32_t in_state[16]);
 void cpu_chacha20_block(void* out_buffer, const void** in_buffers);
diff --git a/lib/defs.hpp b/lib/defs.hpp
index c165a3f..8e49838 100644
--- a/lib/defs.hpp
+++ b/lib/defs.hpp
@@ -8,6 +8,8 @@ typedef unsigned int uint;
 constexpr uint ChaChaDoubleRoundCount = 10;
 constexpr uint ChaChaStateSizeInWords = 16;
 constexpr uint ChaChaStateSizeInBytes = 4 * ChaChaStateSizeInWords;
+constexpr uint ChaChaStateWordsPerRow = 4;
+constexpr uint ChaChaStateSizeInRows = ChaChaStateSizeInWords / ChaChaStateWordsPerRow;
 
 #ifdef CUDA_ENABLED
 // Constants for Cuda kernels
diff --git a/lib/generic/cpu_kernel_arch.hpp b/lib/generic/cpu_kernel_arch.hpp
new file mode 100644
index 0000000..51a63b0
--- /dev/null
+++ b/lib/generic/cpu_kernel_arch.hpp
@@ -0,0 +1,70 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: © 2022 Aalto University
+
+#pragma once
+
+#include "../defs.hpp"
+#include <stdint.h>
+#include <algorithm>
+
+struct StateRow
+{
+    uint32_t values[ChaChaStateWordsPerRow];
+
+    StateRow() : values() { }
+    StateRow(const uint32_t row_values[ChaChaStateWordsPerRow]) : StateRow()
+    {
+        std::copy_n(row_values, ChaChaStateWordsPerRow, values);
+    }
+
+    inline StateRow& operator+=(const StateRow other)
+    {
+        for (uint i = 0; i < ChaChaStateWordsPerRow; ++i)
+        {
+            values[i] += other.values[i];
+        }
+        return *this;
+    }
+
+    inline StateRow& operator^=(const StateRow other)
+    {
+        for (uint i = 0; i < ChaChaStateWordsPerRow; ++i)
+        {
+            values[i] ^= other.values[i];
+        }
+        return *this;
+    }
+
+    inline StateRow& operator<<=(int num_bits)
+    {
+        for (uint i = 0; i < ChaChaStateWordsPerRow; ++i)
+        {
+            uint32_t val = values[i];
+            values[i] = (val << num_bits) | (val >> (32 - num_bits));
+        }
+        return *this;
+    }
+
+    template <int num_positions>
+    inline StateRow rotate_elements_left() const
+    {
+        StateRow res;
+        for (uint i = 0; i < ChaChaStateWordsPerRow; ++i)
+        {
+            res.values[i] = values[(num_positions + i) % ChaChaStateWordsPerRow];
+        }
+        return res;
+    }
+
+    template <uint num_positions>
+    inline StateRow rotate_elements_right() const
+    {
+        return rotate_elements_left<(ChaChaStateWordsPerRow - num_positions) % ChaChaStateWordsPerRow>();
+    }
+
+    inline void unvectorize(uint32_t out_buffer[ChaChaStateWordsPerRow]) const
+    {
+        std::copy_n(values, ChaChaStateWordsPerRow, out_buffer);
+    }
+
+};
diff --git a/lib/intel/cpu_kernel_arch.hpp b/lib/intel/cpu_kernel_arch.hpp
new file mode 100644
index 0000000..7bc31bb
--- /dev/null
+++ b/lib/intel/cpu_kernel_arch.hpp
@@ -0,0 +1,92 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: © 2022 Aalto University
+
+#pragma once
+
+#include "../defs.hpp"
+#include <stdint.h>
+#include <immintrin.h>
+#include <utility>
+
+
+// Wrap _mm_shuffle_epi32 in a template to enforce that rotation_immediate
+// is always known at compile time.
+template <uint8_t rotation_immediate>
+inline __m128i _mm_shuffle_epi32_templated(__m128i val)
+{
+    return _mm_shuffle_epi32(val, rotation_immediate);
+}
+
+// Rotate elements in a 4-vec
+template <uint num_positions>
+inline __m128i rotate_m128i_left(__m128i vec)
+{
+    constexpr uint8_t rotation_lookup[ChaChaStateWordsPerRow] = {
+        0b11100100,
+        0b00111001,
+        0b01001110,
+        0b10010011
+    };
+    constexpr uint8_t rotation_immediate = rotation_lookup[num_positions];
+    // using the templated wrapper for _mm_shuffle_epi32, otherwise
+    // gcc may be confused and decide to not treat rotation_immediate as
+    // compile-time known in debug mode (-O0) for some reason, resulting
+    // in errors from _mm_shuffle_epi32:
+    return _mm_shuffle_epi32_templated<rotation_immediate>(vec);
+}
+
+template <>
+inline __m128i rotate_m128i_left<0>(__m128i vec)
+{
+    return vec;
+}
+
+struct StateRow
+{
+private:
+    __m128i values;
+
+public:
+    StateRow() {}
+    StateRow(const uint32_t row_values[ChaChaStateWordsPerRow])
+        : values(_mm_load_si128(reinterpret_cast<const __m128i*>(row_values))) { }
+    StateRow(__m128i vals) : values(std::move(vals)) { }
+
+    inline StateRow& operator+=(const StateRow other)
+    {
+        values = _mm_add_epi32(values, other.values);
+        return *this;
+    }
+
+    inline StateRow& operator^=(const StateRow other)
+    {
+        values = _mm_xor_si128(values, other.values);
+        return *this;
+    }
+
+    inline StateRow& operator<<=(int num_bits)
+    {
+        values = _mm_xor_si128(
+            _mm_slli_epi32(values, num_bits),
+            _mm_srli_epi32(values, 32 - num_bits)
+        ); // (value << num_bits) ^ (value >> (32 - num_bits));
+        return *this;
+    }
+
+    template <uint num_positions>
+    inline StateRow rotate_elements_left() const
+    {
+        return StateRow(rotate_m128i_left<num_positions>(values));
+    }
+
+    template <uint num_positions>
+    inline StateRow rotate_elements_right() const
+    {
+        return rotate_elements_left<(ChaChaStateWordsPerRow - num_positions) % ChaChaStateWordsPerRow>();
+    }
+
+    inline void unvectorize(uint32_t out_buffer[ChaChaStateWordsPerRow]) const
+    {
+        _mm_store_si128(reinterpret_cast<__m128i*>(out_buffer), values);
+    }
+};
diff --git a/lib/python_bindings.cpp b/lib/python_bindings.cpp
index 688ee94..451901d 100644
--- a/lib/python_bindings.cpp
+++ b/lib/python_bindings.cpp
@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: © 2021 Aalto University
+// SPDX-FileCopyrightText: © 2022 Aalto University
 
 #include <pybind11/pybind11.h>
 
@@ -14,6 +14,15 @@ constexpr bool cuda_supported()
 #endif
 }
 
+constexpr bool openmp_accelerated()
+{
+#ifdef OPENMP_AVAILABLE
+    return true;
+#else
+    return false;
+#endif
+}
+
 PYBIND11_MODULE(native, m)
 {
     m.def("cpu_chacha20_block_factory",
@@ -25,4 +34,5 @@ PYBIND11_MODULE(native, m)
 #endif // CUDA_ENABLED
 
     m.def("cuda_supported", &cuda_supported, "Returns true if CUDA kernels were compiled.");
+    m.def("openmp_accelerated", &openmp_accelerated, "Returns true if CPU kernels are accelerated using OpenMP.");
 }
diff --git a/setup.py b/setup.py
index af65e67..98c6646 100644
--- a/setup.py
+++ b/setup.py
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: © 2021 Aalto University
+# SPDX-FileCopyrightText: © 2022 Aalto University
 
 import setuptools
 from setuptools import Extension
@@ -70,7 +70,7 @@ def build_extension(self, ext):
 
 _jax_version_lower_constraint = ' >= 0.2.12'
 _jax_version_optimistic_upper_constraint = ', <= 2.0.0'
-_jax_version_upper_constraint = ', <= 0.3.15'
+_jax_version_upper_constraint = ', <= 0.3.17'
 
 _version = version_module.VERSION
 if 'JAX_CHACHA_PRNG_BUILD' in os.environ:


Test patch:
diff --git a/.github/workflows/jax_compatibility_tests.yml b/.github/workflows/jax_compatibility_tests.yml
index 6d062ae..35d2891 100644
--- a/.github/workflows/jax_compatibility_tests.yml
+++ b/.github/workflows/jax_compatibility_tests.yml
@@ -18,7 +18,8 @@ jobs:
         jax-version: [
           0.2.12, 0.2.13, 0.2.14, 0.2.15, 0.2.16,
           0.2.17, 0.2.18, 0.2.19, 0.2.20, 0.2.21,
-          0.2.22, 0.2.27, 0.3.1, 0.3.13, 0.3.15
+          0.2.22, 0.2.27, 0.3.1, 0.3.13, 0.3.15,
+          0.3.17
         ]
 
     steps:
diff --git a/.github/workflows/native_code_tests.yml b/.github/workflows/native_code_tests.yml
new file mode 100644
index 0000000..19ad5b9
--- /dev/null
+++ b/.github/workflows/native_code_tests.yml
@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: © 2022 Aalto University
+
+name: Native code tests
+
+on:
+  workflow_dispatch:
+  push:
+    branches: [ master ]
+  pull_request:
+    branches: [ master, "*stable" ]
+
+jobs:
+  test_native_intel:
+    runs-on: ubuntu-22.04
+    strategy:
+      fail-fast: false
+      matrix:
+        force_generic: ['On', 'Off']
+    steps:
+    - uses: actions/checkout@v2
+      with:
+        submodules: true
+    - name: Install environment
+      run: |
+        sudo apt-get update
+        sudo apt-get install gcc g++ cmake python3 python3-dev
+    - name: Build
+      run: |
+        mkdir build
+        cd build
+        cmake -DBUILD_TESTING=On -DFORCE_GENERIC=${{ matrix.force_generic }} -DCMAKE_BUILD_TYPE=Release ..
+        make
+    - name: Run CPU kernel tests
+      run: |
+        cd build
+        ./cpu_kernel_tests
+
+  test_native_arm_cross:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        force_generic: ['On', 'Off']
+    steps:
+    - uses: actions/checkout@v2
+      with:
+        submodules: true
+    - name: Build and run in container
+      uses: addnab/docker-run-action@v3
+      with:
+        image: ubuntu:jammy
+        options: -v ${{ github.workspace }}:/work/src -e FORCE_GENERIC=${{ matrix.force_generic }}
+        run: |
+          sed -i 's/http/[arch=i386,amd64] http/g' /etc/apt/sources.list
+          echo "deb [arch=arm64] http://ports.ubuntu.com/ jammy main restricted" >> /etc/apt/sources.list
+          echo "deb [arch=arm64] http://ports.ubuntu.com/ jammy-updates main restricted" >> /etc/apt/sources.list
+          dpkg --add-architecture arm64
+          apt-get update
+          apt-get upgrade -y
+          apt-get install -y qemu binfmt-support qemu-user-static libc6:arm64 \
+            gcc-aarch64-linux-gnu g++-aarch64-linux-gnu libc6-dev-arm64-cross cmake python3 python3-dev libpython3-dev:arm64
+          mkdir /work/build
+          cd /work/build
+          cmake \
+            -DCMAKE_SYSTEM_NAME=Linux -DCMAKE_SYSTEM_PROCESSOR=aarch64 \
+            -DCMAKE_C_COMPILER=/usr/bin/aarch64-linux-gnu-gcc \
+            -DCMAKE_CXX_COMPILER=/usr/bin/aarch64-linux-gnu-g++ \
+            -DCMAKE_FIND_ROOT_PATH=/usr/aarch64-linux-gnu/ \
+            -DCMAKE_FIND_ROOT_PATH_MODE_PROGRAM=BOTH \
+            -DCMAKE_CROSSCOMPILING_EMULATOR="/usr/bin/qemu-aarch64-static;-L;/usr/aarch64-linux-gnu/" \
+            -DBUILD_TESTING=On \
+            -DFORCE_GENERIC=${{ matrix.force_generic }} \
+            -DCMAKE_BUILD_TYPE=Release ../src
+          make
+          /usr/bin/qemu-aarch64-static -L /usr/aarch64-linux-gnu/ ./cpu_kernel_tests
diff --git a/.github/workflows/python_unittests.yml b/.github/workflows/python_unittests.yml
index 107b276..c3ad8e3 100644
--- a/.github/workflows/python_unittests.yml
+++ b/.github/workflows/python_unittests.yml
@@ -73,7 +73,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v2
       with:
-        python-version: 3.8
+        python-version: "3.10"
     - name: Install environment
       run: |
         python -m pip install --upgrade pip
diff --git a/cmake_config/neon_test.cpp b/cmake_config/neon_test.cpp
new file mode 100644
index 0000000..ef5e7b4
--- /dev/null
+++ b/cmake_config/neon_test.cpp
@@ -0,0 +1,27 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: © 2022 Aalto University
+
+#include <array>
+#include <cstdint>
+#include <arm_neon.h>
+
+int main(int argc, const char** argv)
+{
+    std::array<uint32_t, 4> x { 1, 2, 3, 4 };
+    std::array<uint32_t, 4> y { 5, 6, 7, 8 };
+
+    uint32x4_t x_vec = vld1q_u32(x.data());
+    uint32x4_t y_vec = vld1q_u32(y.data());
+
+    uint32x4_t z_vec = vaddq_u32(x_vec, y_vec);
+
+    std::array<uint32_t, 4> z;
+    vst1q_u32(z.data(), z_vec);
+
+    for (int i = 0; i < 4; ++i)
+    {
+        if (z[i] != x[i] + y[i])
+            return 1;
+    }
+    return 0;
+}
diff --git a/cmake_config/sse_test.cpp b/cmake_config/sse_test.cpp
new file mode 100644
index 0000000..a710741
--- /dev/null
+++ b/cmake_config/sse_test.cpp
@@ -0,0 +1,27 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: © 2022 Aalto University
+
+#include <array>
+#include <cstdint>
+#include <immintrin.h>
+
+int main(int argc, const char** argv)
+{
+    std::array<uint32_t, 4> x { 1, 2, 3, 4 };
+    std::array<uint32_t, 4> y { 5, 6, 7, 8 };
+
+    __m128i x_vec = _mm_load_si128(reinterpret_cast<const __m128i*>(x.data()));
+    __m128i y_vec = _mm_load_si128(reinterpret_cast<const __m128i*>(y.data()));
+
+    __m128i z_vec = _mm_add_epi32(x_vec, y_vec);
+
+    std::array<uint32_t, 4> z;
+    _mm_store_si128(reinterpret_cast<__m128i*>(z.data()), z_vec);
+
+    for (int i = 0; i < 4; ++i)
+    {
+        if (z[i] != x[i] + y[i])
+            return 1;
+    }
+    return 0;
+}
diff --git a/tests/lib/cpu_kernel_tests.cpp b/tests/lib/cpu_kernel_tests.cpp
index 2bdecd1..f8d9054 100644
--- a/tests/lib/cpu_kernel_tests.cpp
+++ b/tests/lib/cpu_kernel_tests.cpp
@@ -6,19 +6,73 @@
 #include "tests.hpp"
 #include "cpu_kernel.hpp"
 
-inline __m128i from_array(std::array<uint32_t, 4>& vals)
+inline StateRow from_array(const std::array<uint32_t, 4>& vals)
 {
-    return _mm_load_si128(reinterpret_cast<__m128i*>(vals.data()));
+    return StateRow(vals.data());
 }
 
-
-inline std::array<uint32_t, 4> to_array(__m128i x)
+inline std::array<uint32_t, 4> to_array(const StateRow& x)
 {
     std::array<uint32_t, 4> vals;
-    _mm_store_si128(reinterpret_cast<__m128i*>(vals.data()), x);
+    x.unvectorize(vals.data());
     return vals;
 }
 
+template <size_t len>
+void print_array(const std::array<uint32_t, len>& arr)
+{
+    for (size_t i = 0; i < len; ++i)
+    {
+        std::cout << arr[i] << ", ";
+    }
+    std::cout << std::endl;
+}
+
+
+template <size_t len>
+void print_array(const uint32_t arr[len])
+{
+    for (size_t i = 0; i < len; ++i)
+    {
+        std::cout << arr[i] << ", ";
+    }
+    std::cout << std::endl;
+}
+
+int test_vectorize_unvectorize()
+{
+    std::cout << "test_vectorize_unvectorize(16 element array): ";
+    int num_fails = 0;
+    std::array<uint32_t, 16> vals({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
+    std::array<uint32_t, 16> expected(vals);
+    std::array<uint32_t, 16> result;
+
+    VectorizedState state(vals.data());
+    state.unvectorize(result.data());
+
+    num_fails += test_assert(result == expected);
+
+    std::cout << "test_vectorize_unvectorize(4 vectors): ";
+
+    std::array<uint32_t, 4> first_row_vals({1, 2, 3, 4});
+    std::array<uint32_t, 4> second_row_vals({5, 6, 7, 8});
+    std::array<uint32_t, 4> third_row_vals({9, 10, 11, 12});
+    std::array<uint32_t, 4> fourth_row_vals({13, 14, 15, 16});
+
+    // auto first_row = from_array(first_row_vals);
+    // auto second_row = from_array(second_row_vals);
+    auto third_row = from_array(third_row_vals);
+    auto fourth_row = from_array(fourth_row_vals);
+
+    state = VectorizedState(
+        from_array(first_row_vals), from_array(second_row_vals), third_row, fourth_row
+    );
+    state.unvectorize(result.data());
+    num_fails += test_assert(result == expected);
+
+    return num_fails;
+}
+
 int test_rotate_elements_left()
 {
     int num_fails = 0;
@@ -26,69 +80,157 @@ int test_rotate_elements_left()
     std::array<uint32_t, 4> expected;
     std::array<uint32_t, 4> result;
 
-    auto vec_vals = from_array(vals);
+    auto vals_vec = from_array(vals);
 
     std::cout << "test_rotate_elements_left(0): ";
     expected = vals;
-    result = to_array(rotate_elements_left<0>(vec_vals));
+    result = to_array(vals_vec.rotate_elements_left<0>());
     num_fails += test_assert(result == expected);
 
     std::cout << "test_rotate_elements_left(1): ";
     expected = {2, 3, 4, 1};
-    result = to_array(rotate_elements_left<1>(vec_vals));
+    result = to_array(vals_vec.rotate_elements_left<1>());
+    num_fails += test_assert(result == expected);
+
+    std::cout << "test_rotate_elements_left(2): ";
+    expected = {3, 4, 1, 2};
+    result = to_array(vals_vec.rotate_elements_left<2>());
     num_fails += test_assert(result == expected);
 
     std::cout << "test_rotate_elements_left(3): ";
     expected = {4, 1, 2, 3};
-    result = to_array(rotate_elements_left<3>(vec_vals));
+    result = to_array(vals_vec.rotate_elements_left<3>());
+    num_fails += test_assert(result == expected);
+
+    return num_fails;
+}
+
+int test_rotate_left()
+{
+    int num_fails = 0;
+
+    std::array<uint32_t, 4> vals({ 0x80000001, 0xffffffff, 0x0f0f0f0f, 0xcccccccc });
+    auto vec_vals = from_array(vals);
+
+    std::cout << "test_rotate_left(x, 1): ";
+    std::array<uint32_t, 4> expected = { 3, 0xffffffff, 0x1e1e1e1e, 0x99999999 };
+
+    auto vec_result = vec_vals << 1;
+    auto result = to_array(vec_result);
+
+    num_fails += test_assert(result == expected);
+
+    std::cout << "test_rotate_left(x, 15): ";
+    expected = { 0x0000c000, 0xffffffff, 0x87878787, 0x66666666 };
+
+    vec_result = vec_vals << 15;
+    result = to_array(vec_result);
+
+    num_fails += test_assert(result == expected);
+
+    std::cout << "test_rotate_left(x, 16): ";
+    expected = { 0x00018000, 0xffffffff, 0x0f0f0f0f, 0xcccccccc };
+
+    vec_result = vec_vals << 16;
+    result = to_array(vec_result);
+
+    num_fails += test_assert(result == expected);
+
+    std::cout << "test_rotate_left(x, 31): ";
+    expected = { 0xc0000000, 0xffffffff, 0x87878787, 0x66666666 };
+
+    vec_result = vec_vals << 31;
+    result = to_array(vec_result);
+
+    num_fails += test_assert(result == expected);
+
+    return num_fails;
+}
+
+int test_add()
+{
+    int num_fails = 0;
+
+    std::array<uint32_t, 4> left_vals({ 1, 0xffffffff, 0x0f0f0f0f, 0xcccccccc });
+    std::array<uint32_t, 4> right_vals({ 1, 1, 0x0f0f0f0f, 0x00cc00cc });
+
+    std::array<uint32_t, 4> expected({ 2, 0, 0x1e1e1e1e , 0xcd98cd98});
+
+    std::cout << "test_add(): ";
+    auto left = from_array(left_vals);
+    auto right = from_array(right_vals);
+    auto vec_result = left + right;
+    auto result = to_array(vec_result);
+
     num_fails += test_assert(result == expected);
 
     return num_fails;
 }
 
-// int test_rotate_left()
-// {
-//     int num_fails = 0;
+int test_xor()
+{
+    int num_fails = 0;
 
-//     std::array<uint32_t, 4> vals({ 1, 0xffffffff, 0x0f0f0f0f, 0xcccccccc });
-//     __m128i vec_vals = from_array(vals);
+    std::array<uint32_t, 4> left_vals({ 1, 0xffffffff, 0x0f0f0f0f, 0xcccccccc });
+    std::array<uint32_t, 4> right_vals({ 1, 1, 0xf0f0f0f0, 0x00cc00cc });
 
-//     std::cout << "test_rotate_left(x, 0):\t\t";
-//     std::array<uint32_t, 4> expected(vals);
+    std::array<uint32_t, 4> expected({ 0, 0xfffffffe, 0xffffffff , 0xcc00cc00 });
 
-//     __m128i vec_result = rotate_left(vec_vals, 0);
-//     std::array<uint32_t, 4> result = to_array<uint32_t>(vec_result);
+    std::cout << "test_xor(): ";
+    auto left = from_array(left_vals);
+    auto right = from_array(right_vals);
+    auto vec_result = left ^ right;
+    auto result = to_array(vec_result);
 
-//     num_fails += test_assert(vals == result);
+    num_fails += test_assert(result == expected);
 
-//     std::cout << "test_rotate_left(x, 1):\t\t";
-//     expected = { 2, 0xffffffff, 0x1e1e1e1e, 0x99999999 };
+    return num_fails;
+}
 
-//     vec_result = rotate_left(vec_vals, 1);
-//     result = to_array<uint32_t>(vec_result);
+int test_pack_diagonals()
+{
+    int num_fails = 0;
 
-//     num_fails += test_assert(result == expected);
+    std::cout << "test_pack_diagonals(): ";
+    std::array<uint32_t, 16> vals({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
+    std::array<uint32_t, 16> expected({1, 2, 3, 4, 6, 7, 8, 5, 11, 12, 9, 10, 16, 13, 14, 15});
+    std::array<uint32_t, 16> result;
 
-//     std::cout << "test_rotate_left(x, 32):\t";
-//     expected = vals;
+    auto vec = VectorizedState(vals.data());
+    pack_diagonals(vec, vec);
+    vec.unvectorize(result.data());
+    num_fails += test_assert(result == expected);
 
-//     vec_result = rotate_left(vec_vals, 32);
-//     result = to_array<uint32_t>(vec_result);
+    return num_fails;
+}
+
+
+int test_unpack_diagonals()
+{
+    int num_fails = 0;
 
-//     num_fails += test_assert(result == expected);
+    std::cout << "test_unpack_diagonals(): ";
+    std::array<uint32_t, 16> expected({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
+    std::array<uint32_t, 16> vals({1, 2, 3, 4, 6, 7, 8, 5, 11, 12, 9, 10, 16, 13, 14, 15});
+    std::array<uint32_t, 16> result;
 
-//     return num_fails;
-// }
+    auto vec = VectorizedState(vals.data());
+    unpack_diagonals(vec, vec);
+    vec.unvectorize(result.data());
+    num_fails += test_assert(result == expected);
 
-int test_cpu_chacha20_block_sse()
+    return num_fails;
+}
+
+int test_chacha20_block()
 {
     int num_fails = 0;
     for (int i = 0; i < test_vector_states.size(); ++i)
     {
         std::array<uint32_t, ChaChaStateSizeInWords> result;
 
-        std::cout << "test_chacha20_block_sse(vector " << i << "): ";
-        chacha20_block_sse(result.data(), test_vector_states[i].data());
+        std::cout << "test_chacha20_block(vector " << i << "): ";
+        chacha20_block(result.data(), test_vector_states[i].data());
         num_fails += test_assert(result == test_vector_expected[i]);
     }
 
@@ -184,8 +326,14 @@ int main(int argc, const char** argv)
 {
     int num_fails = 0;
     // test_rotate_left();
+    num_fails += test_vectorize_unvectorize();
     num_fails += test_rotate_elements_left();
-    num_fails += test_cpu_chacha20_block_sse();
+    num_fails += test_rotate_left();
+    num_fails += test_add();
+    num_fails += test_xor();
+    num_fails += test_pack_diagonals();
+    num_fails += test_unpack_diagonals();
+    num_fails += test_cpu_chacha20_block();
     num_fails += test_cpu_chacha20_block();
     return (num_fails > 0);
 }
diff --git a/tests/random_test.py b/tests/random_test.py
index d7ae627..00fa5fc 100644
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -7,7 +7,7 @@
 import jax
 
 from chacha.random import split, PRNGKey, random_bits, uniform, _uniform, is_state_invalidated, ErrorFlag
-from chacha.cipher import set_counter
+from chacha.cipher import increase_counter, set_counter, get_counter, increment_counter
 import numpy as np  # type: ignore
 
 
@@ -22,26 +22,34 @@ def test_random_bits(self) -> None:
         ], dtype=jnp.uint32)
 
         shape = (17, 9)
-        x, error = random_bits(rng_key, 8, shape)
+        x, next_rng_key, error = random_bits(rng_key, 8, shape)
         self.assertEqual(0, error)
         self.assertEqual(x.dtype, jnp.uint8)
         self.assertEqual(x.shape, shape)
+        expected_next_rng_key = increase_counter(rng_key, 3)
+        self.assertTrue(np.all(expected_next_rng_key == next_rng_key))
 
-        x, error = random_bits(rng_key, 16, shape)
+        x, next_rng_key, error = random_bits(rng_key, 16, shape)
         self.assertEqual(0, error)
         self.assertEqual(x.dtype, jnp.uint16)
         self.assertEqual(x.shape, shape)
+        expected_next_rng_key = increase_counter(rng_key, 5)
+        self.assertTrue(np.all(expected_next_rng_key == next_rng_key))
 
-        x, error = random_bits(rng_key, 32, shape)
+        x, next_rng_key, error = random_bits(rng_key, 32, shape)
         self.assertEqual(0, error)
         self.assertEqual(x.dtype, jnp.uint32)
         self.assertEqual(x.shape, shape)
+        expected_next_rng_key = increase_counter(rng_key, 10)
+        self.assertTrue(np.all(expected_next_rng_key == next_rng_key))
 
         # needs 64bit jax mode
-        x, error = random_bits(rng_key, 64, shape)
+        x, next_rng_key, error = random_bits(rng_key, 64, shape)
         self.assertEqual(0, error)
         self.assertEqual(x.dtype, jnp.uint64)
         self.assertEqual(x.shape, shape)
+        expected_next_rng_key = increase_counter(rng_key, 20)
+        self.assertTrue(np.all(expected_next_rng_key == next_rng_key))
 
     def test_is_state_invalidated(self) -> None:
         rng_key = jnp.array([
@@ -67,7 +75,7 @@ def test_split_from_initial(self) -> None:
             [0x61707865, 0x3320646e, 0x79622d32, 0x6b206574],
             [0x00000000, 0x00000000, 0x00000000, 0x00000000],
             [0x00000000, 0x00000000, 0x00000000, 0x00000000],
-            [0xe56a5d40, 0x00000000, 0x00000000, 0x00000001],
+            [0x00000000, 0x00000000, 0x00000000, 0x00000001],
         ], dtype=jnp.uint32)
 
         # 6 splits from initial
@@ -153,7 +161,7 @@ def test_split_invalidates_when_split_limit_exceeded(self) -> None:
             [0x61707865, 0x3320646e, 0x79622d32, 0x6b206574],
             [0xbbbbbbbb, 0x00000000, 0x00000000, 0x00000000],
             [0x00000000, 0x00000000, 0x00000000, 0xaaaaaaaa],
-            [0xbe57823a, 0x40000000, 0x00000000, 0x00000000],
+            [0x00000000, 0x40000000, 0x00000000, 0x00000000],
         ], dtype=jnp.uint32)
 
         split_keys = split(rng_key, 5)
@@ -166,7 +174,7 @@ def test_split_reproduces_invalid_states(self) -> None:
             [0x61707865, 0x3320646e, 0x79622d32, 0x6b206574],
             [0xbbbbbbbb, 0x00000000, 0x00000000, 0x00000000],
             [0x00000000, 0x00000000, 0x00000000, 0xaaaaaaaa],
-            [0xbe57823a, 0x00000000, 0x00000000, 0x00000000],
+            [0x00000000, 0x00000000, 0x00000000, 0x00000000],
         ], dtype=jnp.uint32)
         assert is_state_invalidated(rng_key)
 
@@ -180,6 +188,19 @@ def test_split_rejects_too_large_split(self) -> None:
         with self.assertRaises(ValueError):
             split(rng_key, 2**32)
 
+    def test_split_invalidates_when_nonzero_counter(self) -> None:
+        rng_key = jnp.array([
+            [0x61707865, 0x3320646e, 0x79622d32, 0x6b206574],
+            [0xbbbbbbbb, 0x00000000, 0x00000000, 0x00000000],
+            [0x00000000, 0x00000000, 0x00000000, 0xaaaaaaaa],
+            [0xbe57823a, 0x00000000, 0x00000000, 0x00000006],
+        ], dtype=jnp.uint32)
+
+        split_keys = split(rng_key, 5)
+
+        for split_key in split_keys:
+            self.assertTrue(is_state_invalidated(split_key))
+
     def test_PRNGKey_from_int(self) -> None:
         x = 9651354789628635673475235
 
@@ -236,6 +257,21 @@ def test_uniform(self) -> None:
         self.assertTrue(np.all(data >= minval))
         self.assertTrue(np.all(data <= maxval))
 
+    def test_uniform_with_next_key(self) -> None:
+        rng_key = PRNGKey(0)
+        shape = (3, 12)
+        minval = 1
+        maxval = 3
+        dtype = jnp.float32
+        data, next_rng_key = uniform(rng_key, shape, dtype, minval, maxval, return_next_key=True)
+
+        self.assertEqual(jnp.float32, data.dtype)
+        self.assertEqual(shape, data.shape)
+        self.assertTrue(np.all(data >= minval))
+        self.assertTrue(np.all(data <= maxval))
+        expected_next_rng_key = increase_counter(rng_key, 3)
+        self.assertTrue(np.all(expected_next_rng_key == next_rng_key))
+
     def test_uniform_scalar(self) -> None:
         rng_key = PRNGKey(0)
         sample = uniform(rng_key)
@@ -255,6 +291,26 @@ def test_uniform_vmap(self) -> None:
         self.assertTrue(np.all(data >= minval))
         self.assertTrue(np.all(data <= maxval))
 
+    def test_uniform_vmap_with_next_key(self) -> None:
+        rng_key = PRNGKey(0)
+        batch_keys = split(rng_key, 7)
+        shape = (3, 12)
+        minval = 1
+        maxval = 3
+        dtype = jnp.float32
+        data, next_rng_keys = jax.vmap(
+            lambda key: uniform(key, shape, dtype, minval, maxval, return_next_key=True)
+        )(batch_keys)
+
+        self.assertEqual(jnp.float32, data.dtype)
+        self.assertEqual((7, *shape), data.shape)
+        self.assertTrue(np.all(data >= minval))
+        self.assertTrue(np.all(data <= maxval))
+        expected_next_rng_keys = np.zeros_like(batch_keys, dtype=batch_keys.dtype)
+        for i in range(7):
+            expected_next_rng_keys[i] = increase_counter(batch_keys[i], 3)
+        self.assertTrue(np.all(expected_next_rng_keys == next_rng_keys))
+
     def test_uniform_scalar_vmap(self) -> None:
         rng_key = PRNGKey(0)
         batch_keys = split(rng_key, 7)
@@ -265,19 +321,26 @@ def test_random_bits_consistency(self) -> None:
         """ verifies that there is no difference between sampling two blocks directly
             or separately (while manually increasing counter) """
         rng_key = PRNGKey(0)
-        single_vals, _ = random_bits(rng_key, 32, (32,))
-        multi_vals_1, _ = random_bits(rng_key, 32, (16,))
-        rng_key_2 = set_counter(rng_key, 1)
-        multi_vals_2, _ = random_bits(rng_key_2, 32, (16,))
+        single_vals, next_rng_key, _ = random_bits(rng_key, 32, (32,))
+        multi_vals_1, next_rng_key_multi_1, _ = random_bits(rng_key, 32, (16,))
+        multi_vals_2, next_rng_key_multi_2, _ = random_bits(next_rng_key_multi_1, 32, (16,))
         self.assertTrue(np.all(single_vals[:16] == multi_vals_1))
         self.assertTrue(np.all(single_vals[16:] == multi_vals_2))
 
+        expected_next_rng_key = increase_counter(rng_key, 2)
+        self.assertTrue(np.all(expected_next_rng_key == next_rng_key))
+        expected_next_rng_key_multi_1 = increase_counter(rng_key, 1)
+        self.assertTrue(np.all(expected_next_rng_key_multi_1 == next_rng_key_multi_1))
+        self.assertTrue(np.all(expected_next_rng_key == next_rng_key_multi_2))
+
     def test_random_bits_scalar(self) -> None:
         rng_key = PRNGKey(0)
-        sample, error = random_bits(rng_key, 32, ())
+        sample, next_rng_key, error = random_bits(rng_key, 32, ())
         self.assertEqual(0, error)
         self.assertEqual((), jnp.shape(sample))
         self.assertEqual(jnp.uint32, jnp.dtype(sample))
+        expected_next_rng_key = increase_counter(rng_key, 1)
+        self.assertTrue(np.all(expected_next_rng_key == next_rng_key))
 
     def test_random_bits_invalid_width(self) -> None:
         rng_key = PRNGKey(0)
@@ -288,7 +351,7 @@ def test_uniform_invalid_dtype(self) -> None:
         rng_key = PRNGKey(0)
 
         with self.assertRaises(TypeError):
-            _uniform(rng_key, (), jnp.uint8, 0., 1.)
+            _uniform(rng_key, (), jnp.uint8, 0., 1., False)
         with self.assertRaises(TypeError):
             uniform(rng_key, (), jnp.uint32)
 
@@ -300,9 +363,10 @@ def test_random_bits_invalid_state(self) -> None:
             [0xbe57823a, 0x00000000, 0x00000000, 0x00000000],
         ], dtype=jnp.uint32)
 
-        x, error = random_bits(rng_key, 32, (1, 5))
+        x, next_rng_key, error = random_bits(rng_key, 32, (1, 5))
         self.assertTrue(jnp.all(x == 0))
         self.assertEqual(ErrorFlag.InvalidState, error)
+        self.assertTrue(is_state_invalidated(next_rng_key))
 
     def test_random_bits_counter_overflow(self) -> None:
         rng_key = jnp.array([
@@ -313,9 +377,10 @@ def test_random_bits_counter_overflow(self) -> None:
         ], dtype=jnp.uint32)
         # state has one block left to generate
 
-        x, error = random_bits(rng_key, 32, (4, 5))  # will require generating two blocks
+        x, next_rng_key, error = random_bits(rng_key, 32, (4, 5))  # will require generating two blocks
         self.assertTrue(jnp.all(x == 0))
         self.assertEqual(ErrorFlag.CounterOverflow, error)
+        self.assertTrue(is_state_invalidated(next_rng_key))
 
     def test_uniform_invalid_state(self) -> None:
         rng_key = jnp.array([
@@ -340,6 +405,31 @@ def test_uniform_counter_overflow(self) -> None:
         x = uniform(rng_key, (3, 2, 1))
         self.assertTrue(jnp.all(jnp.isnan(x)))
 
+    def test_uniform_invalid_state_with_next_key(self) -> None:
+        rng_key = jnp.array([
+            [0x61707865, 0x3320646e, 0x79622d32, 0x6b206574],
+            [0xbbbbbbbb, 0x00000000, 0x00000000, 0x00000000],
+            [0x00000000, 0x00000000, 0x00000000, 0xaaaaaaaa],
+            [0xbe57823a, 0x00000000, 0x00000000, 0x00000000],
+        ], dtype=jnp.uint32)
+
+        x, next_rng_key = uniform(rng_key, (3, 2, 1), return_next_key=True)
+        self.assertTrue(jnp.all(jnp.isnan(x)))
+        self.assertTrue(is_state_invalidated(next_rng_key))
+
+    def test_uniform_counter_overflow_with_next_key(self) -> None:
+        rng_key = jnp.array([
+            [0x61707865, 0x3320646e, 0x79622d32, 0x6b206574],
+            [0xbbbbbbbb, 0x00000000, 0x00000000, 0x00000000],
+            [0x00000000, 0x00000000, 0x00000000, 0xaaaaaaaa],
+            [0xffffffff, 0x00000000, 0x00000000, 0x00000001],
+        ], dtype=jnp.uint32)
+        # state has zero blocks left to generate
+
+        x, next_rng_key = uniform(rng_key, (3, 2, 1), return_next_key=True)
+        self.assertTrue(jnp.all(jnp.isnan(x)))
+        self.assertTrue(is_state_invalidated(next_rng_key))
+
 
 if __name__ == '__main__':
     unittest.main()