Which XlaOps are used for random number generation in JAX?

[Ebanflo42]

I created an XLA issue to try to understand this here.

I'm interested in understanding how a high-level API can be built on top of XLA in general, and RNGs are essential to this. In summary, we have the op XlaOp RngBitGenerator(RandomAlgorithm algorithm, XlaOp initial_state, const Shape& shape);, which takes a seed but isn't a specific distribution, whereas Ops RngUniform and RngNormal are specific distributions that apparently don't take a seed.

In contrast, in JAX, every RNG function takes a specific seed and generates a specific distribution. I'd like to understand how this works at a lower level and I'm not sure where to look.

[jaro-sevcik]

The random generators and distributions are implement completely in JAX, see the jax/_src/random.py file. For example, see here for the code for jax.random.normal.

If you are interested in seeing the XLA code for random number generation, you can dump the MLIR code resulting from JIT compilation of the particular random function. Example:

import jax

key = jax.random.key(0)
print(jax.jit(jax.random.normal).lower(key).as_text())

This will produce something like this (note that the body of the _normal_real function indeed corresponds to the JAX code for _normal_real):

module @jit_normal attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<2xui32> {mhlo.layout_mode = "default"}) -> (tensor<f32> {jax.result_info = "", mhlo.layout_mode = "default"}) {
    %0 = stablehlo.custom_call @Sharding(%arg0) {mhlo.sharding = "{replicated}"} : (tensor<2xui32>) -> tensor<2xui32>
    %1 = call @_normal(%0) : (tensor<2xui32>) -> tensor<f32>
    return %1 : tensor<f32>
  }
  func.func private @_normal(%arg0: tensor<2xui32>) -> tensor<f32> {
    %0 = call @_normal_real(%arg0) : (tensor<2xui32>) -> tensor<f32>
    return %0 : tensor<f32>
  }
  func.func private @_normal_real(%arg0: tensor<2xui32>) -> tensor<f32> {
    %0 = stablehlo.constant dense<-0.99999994> : tensor<f32>
    %1 = stablehlo.constant dense<1.000000e+00> : tensor<f32>
    %2 = call @_uniform(%arg0, %0, %1) : (tensor<2xui32>, tensor<f32>, tensor<f32>) -> tensor<f32>
    %3 = chlo.erf_inv %2 : tensor<f32> -> tensor<f32>
    %4 = stablehlo.constant dense<1.41421354> : tensor<f32>
    %5 = stablehlo.multiply %4, %3 : tensor<f32>
    return %5 : tensor<f32>
  }
  func.func private @_uniform(%arg0: tensor<2xui32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
... SNIP ...