Repository: ml-explore/mlx-swift
Pull request #219: various Random improvements
Author: login: davidkoski
Created: 2025-04-08T19:47:33Z | Merged: 2025-04-18T16:54:15Z

- allow a Key or a RandomState to be used with Random functions
- users can create RandomState
- users can do withRandomState to scope the use of a RandomState -- useful with multiple threads


Diff:
diff --git a/Source/MLX/Random.swift b/Source/MLX/Random.swift
index a42fcfaf..2e7a67e6 100644
--- a/Source/MLX/Random.swift
+++ b/Source/MLX/Random.swift
@@ -3,6 +3,87 @@
 import Cmlx
 import Foundation
 
+/// Collection of functions related to random number generation.
+///
+/// Following [JAX’s PRNG design](https://jax.readthedocs.io/en/latest/jep/263-prng.html)
+/// we use a splittable version of Threefry, which is a counter-based PRNG.
+///
+/// Random sampling functions in MLX use an implicit global PRNG state by default.
+/// However, all functions take an optional key keyword argument for when more fine-grained
+/// control or explicit state management is needed.  Callers can also arrange for `Task` local
+/// random state -- useful in multithreaded situations.
+///
+/// For example, you can generate random numbers with:
+///
+/// ```swift
+/// for _ in 0 ..< 3 {
+///   print(MLXRandom.uniform())
+/// }
+/// ```
+///
+/// which will print a sequence of unique pseudo random numbers. Alternatively you can explicitly set the key:
+///
+/// ```swift
+/// let key = MLXRandom.key(0)
+/// for _ in 0 ..< 3 {
+///   print(MLXRandom.uniform(key: key))
+/// }
+/// ```
+///
+/// which will yield the same pseudo random number at each iteration as the key doesn't change.
+///
+/// To get a new random number for each call you would ``split(key:stream:)`` the key:
+///
+/// ```swift
+/// var key = MLXRandom.key(0)
+/// for _ in 0 ..< 3 {
+///   let (a, b) = MLXRandom.split(key: key)
+///
+///   // use b to generate a different value each time
+///   print(MLXRandom.uniform(key: b))
+///
+///   // new random state is a
+///   key = a
+/// }
+/// ```
+///
+/// This will generate the same sequence of numbers each time (same starting key) but
+/// different values for each call.
+///
+/// As a convenience you can use ``RandomState`` to manage the key splitting and:
+///
+/// ```swift
+/// let state = RandomState(seed: 0)
+/// for _ in 0 ..< 3 {
+///   print(MLXRandom.uniform(key: state))
+/// }
+/// ```
+///
+/// Finally, if you need to control random state in deeply nested calls to `MLXRandom` or you need
+/// thread-safe random state for multi-threaded evaluation you can use ``withRandomState(_:body:)-6i2p1``:
+///
+/// ```swift
+/// await withTaskGroup { group in
+///     for i in 0 ..< 10 {
+///         group.addTask {
+///             let state = MLXRandom.RandomState(seed: UInt64(i))
+///             return withRandomState(state) {
+///                 var t: Float = 0.0
+///                 for _ in 0 ..< 100 {
+///                     t += uniform(0 ..< 1, [10, 10]).sum().item(Float.self)
+///                 }
+///                 return t
+///             }
+///         }
+///     }
+///
+///     for await v in group {
+///         ...
+///     }
+/// }
+/// ```
+///
+/// Each task will have separate ``RandomState`` that will be used implicitly (if no other key is passed in).
 public enum MLXRandom {
 
     /// Seed the global PRNG.
@@ -67,12 +148,13 @@ public enum MLXRandom {
     /// let array = MLXRandom.uniform(0.5 ..< 1, [50], key: key)
     /// ```
     public static func uniform<R: HasDType, T>(
-        _ range: Range<R>, _ shape: [Int] = [], type: T.Type = Float.self, key: MLXArray? = nil,
+        _ range: Range<R>, _ shape: [Int] = [], type: T.Type = Float.self,
+        key: RandomStateOrKey? = nil,
         stream: StreamOrDevice = .default
     ) -> MLXArray where T: HasDType, T: BinaryFloatingPoint {
         let lb = MLXArray(range.lowerBound)
         let ub = MLXArray(range.upperBound)
-        let key = key ?? globalState.next()
+        let key = resolve(key: key)
         var result = mlx_array_new()
 
         mlx_random_uniform(
@@ -92,11 +174,11 @@ public enum MLXRandom {
     /// ```
     public static func uniform<T>(
         _ range: Range<Float> = 0 ..< 1, _ shape: [Int] = [], type: T.Type = Float.self,
-        key: MLXArray? = nil, stream: StreamOrDevice = .default
+        key: RandomStateOrKey? = nil, stream: StreamOrDevice = .default
     ) -> MLXArray where T: HasDType, T: BinaryFloatingPoint {
         let lb = MLXArray(range.lowerBound)
         let ub = MLXArray(range.upperBound)
-        let key = key ?? globalState.next()
+        let key = resolve(key: key)
         var result = mlx_array_new()
 
         mlx_random_uniform(
@@ -122,11 +204,11 @@ public enum MLXRandom {
     /// ```
     public static func uniform<T>(
         low: ScalarOrArray, high: ScalarOrArray, _ shape: [Int]? = nil, type: T.Type = Float.self,
-        key: MLXArray? = nil, stream: StreamOrDevice = .default
+        key: RandomStateOrKey? = nil, stream: StreamOrDevice = .default
     ) -> MLXArray where T: HasDType, T: BinaryFloatingPoint {
         let (low, high) = toArrays(low, high)
         let shape = shape ?? low.shape
-        let key = key ?? globalState.next()
+        let key = resolve(key: key)
         var result = mlx_array_new()
 
         mlx_random_uniform(
@@ -152,11 +234,11 @@ public enum MLXRandom {
     /// ```
     public static func uniform(
         low: ScalarOrArray, high: ScalarOrArray, _ shape: [Int]? = nil, dtype: DType = .float32,
-        key: MLXArray? = nil, stream: StreamOrDevice = .default
+        key: RandomStateOrKey? = nil, stream: StreamOrDevice = .default
     ) -> MLXArray {
         let (low, high) = toArrays(low, high)
         let shape = shape ?? low.shape
-        let key = key ?? globalState.next()
+        let key = resolve(key: key)
         var result = mlx_array_new()
 
         mlx_random_uniform(
@@ -190,10 +272,10 @@ public enum MLXRandom {
     ///   - key: PRNG key
     public static func normal<T>(
         _ shape: [Int] = [], type: T.Type = Float.self, loc: Float = 0, scale: Float = 1,
-        key: MLXArray? = nil,
+        key: RandomStateOrKey? = nil,
         stream: StreamOrDevice = .default
     ) -> MLXArray where T: HasDType, T: BinaryFloatingPoint {
-        let key = key ?? globalState.next()
+        let key = resolve(key: key)
         var result = mlx_array_new()
 
         mlx_random_normal(
@@ -225,10 +307,10 @@ public enum MLXRandom {
     ///   - key: PRNG key
     public static func normal(
         _ shape: [Int] = [], dtype: DType = .float32, loc: Float = 0, scale: Float = 1,
-        key: MLXArray? = nil,
+        key: RandomStateOrKey? = nil,
         stream: StreamOrDevice = .default
     ) -> MLXArray {
-        let key = key ?? globalState.next()
+        let key = resolve(key: key)
         var result = mlx_array_new()
 
         mlx_random_normal(
@@ -255,9 +337,9 @@ public enum MLXRandom {
     ///   - key: PRNG key
     public static func multivariateNormal(
         mean: MLXArray, covariance: MLXArray, shape: [Int] = [], dtype: DType = .float32,
-        key: MLXArray? = nil, stream: StreamOrDevice = .default
+        key: RandomStateOrKey? = nil, stream: StreamOrDevice = .default
     ) -> MLXArray {
-        let key = key ?? globalState.next()
+        let key = resolve(key: key)
         var result = mlx_array_new()
 
         mlx_random_multivariate_normal(
@@ -283,12 +365,12 @@ public enum MLXRandom {
     /// let array = MLXRandom.randInt(Int32(0) ..< 100, [50], key: key)
     /// ```
     public static func randInt<T>(
-        _ range: Range<T>, _ shape: [Int] = [], key: MLXArray? = nil,
+        _ range: Range<T>, _ shape: [Int] = [], key: RandomStateOrKey? = nil,
         stream: StreamOrDevice = .default
     ) -> MLXArray where T: HasDType, T: BinaryInteger {
         let lb = MLXArray(range.lowerBound)
         let ub = MLXArray(range.upperBound)
-        let key = key ?? globalState.next()
+        let key = resolve(key: key)
         var result = mlx_array_new()
 
         mlx_random_randint(
@@ -312,12 +394,13 @@ public enum MLXRandom {
     /// let array = MLXRandom.randInt(low: [0, 10], high: [10, 100], key: key)
     /// ```
     public static func randInt(
-        low: ScalarOrArray, high: ScalarOrArray, _ shape: [Int]? = nil, key: MLXArray? = nil,
+        low: ScalarOrArray, high: ScalarOrArray, _ shape: [Int]? = nil,
+        key: RandomStateOrKey? = nil,
         stream: StreamOrDevice = .default
     ) -> MLXArray {
         let (low, high) = toArrays(low, high)
         let shape = shape ?? low.shape
-        let key = key ?? globalState.next()
+        let key = resolve(key: key)
         var result = mlx_array_new()
 
         mlx_random_randint(
@@ -343,11 +426,11 @@ public enum MLXRandom {
     /// ```
     public static func randInt<T>(
         low: ScalarOrArray, high: ScalarOrArray, _ shape: [Int]? = nil, type: T.Type,
-        key: MLXArray? = nil, stream: StreamOrDevice = .default
+        key: RandomStateOrKey? = nil, stream: StreamOrDevice = .default
     ) -> MLXArray where T: HasDType, T: BinaryInteger {
         let (low, high) = toArrays(low, high)
         let shape = shape ?? low.shape
-        let key = key ?? globalState.next()
+        let key = resolve(key: key)
         var result = mlx_array_new()
 
         mlx_random_randint(
@@ -372,12 +455,12 @@ public enum MLXRandom {
     /// let array = MLXRandom.bernoulli([50, 2], key: key)
     /// ```
     public static func bernoulli(
-        _ shape: [Int] = [], key: MLXArray? = nil, stream: StreamOrDevice = .default
+        _ shape: [Int] = [], key: RandomStateOrKey? = nil, stream: StreamOrDevice = .default
     )
         -> MLXArray
     {
         let p = MLXArray(0.5)
-        let key = key ?? globalState.next()
+        let key = resolve(key: key)
         var result = mlx_array_new()
         mlx_random_bernoulli(&result, p.ctx, shape.asInt32, shape.count, key.ctx, stream.ctx)
 
@@ -403,12 +486,12 @@ public enum MLXRandom {
     /// let array = MLXRandom.bernoulli(MLXArray(convert: [0.1, 0.5, 0.8]), key: key)
     /// ```
     public static func bernoulli(
-        _ p: ScalarOrArray, _ shape: [Int]? = nil, key: MLXArray? = nil,
+        _ p: ScalarOrArray, _ shape: [Int]? = nil, key: RandomStateOrKey? = nil,
         stream: StreamOrDevice = .default
     ) -> MLXArray {
         let p = p.asMLXArray(dtype: .float32)
         let shape = shape ?? p.shape
-        let key = key ?? globalState.next()
+        let key = resolve(key: key)
         var result = mlx_array_new()
         mlx_random_bernoulli(&result, p.ctx, shape.asInt32, shape.count, key.ctx, stream.ctx)
 
@@ -434,12 +517,13 @@ public enum MLXRandom {
     /// ### See also
     /// - [JAX Documentation](https://jax.readthedocs.io/en/latest/_modules/jax/_src/random.html#truncated_normal)
     public static func truncatedNormal<R: HasDType, T>(
-        _ range: Range<R>, _ shape: [Int] = [], type: T.Type = Float.self, key: MLXArray? = nil,
+        _ range: Range<R>, _ shape: [Int] = [], type: T.Type = Float.self,
+        key: RandomStateOrKey? = nil,
         stream: StreamOrDevice = .default
     ) -> MLXArray where T: HasDType, T: BinaryFloatingPoint {
         let lb = MLXArray(range.lowerBound)
         let ub = MLXArray(range.upperBound)
-        let key = key ?? globalState.next()
+        let key = resolve(key: key)
         var result = mlx_array_new()
 
         mlx_random_truncated_normal(
@@ -458,12 +542,13 @@ public enum MLXRandom {
     /// let array = MLXRandom.truncatedNormal(0.5 ..< 1, [50], key: key)
     /// ```
     public static func truncatedNormal<T>(
-        _ range: Range<Float>, _ shape: [Int] = [], type: T.Type = Float.self, key: MLXArray? = nil,
+        _ range: Range<Float>, _ shape: [Int] = [], type: T.Type = Float.self,
+        key: RandomStateOrKey? = nil,
         stream: StreamOrDevice = .default
     ) -> MLXArray where T: HasDType, T: BinaryFloatingPoint {
         let lb = MLXArray(range.lowerBound)
         let ub = MLXArray(range.upperBound)
-        let key = key ?? globalState.next()
+        let key = resolve(key: key)
         var result = mlx_array_new()
 
         mlx_random_truncated_normal(
@@ -488,11 +573,11 @@ public enum MLXRandom {
     /// ```
     public static func truncatedNormal<T>(
         low: ScalarOrArray, high: ScalarOrArray, _ shape: [Int]? = nil, type: T.Type = Float.self,
-        key: MLXArray? = nil, stream: StreamOrDevice = .default
+        key: RandomStateOrKey? = nil, stream: StreamOrDevice = .default
     ) -> MLXArray where T: HasDType, T: BinaryFloatingPoint {
         let (low, high) = toArrays(low, high)
         let shape = shape ?? low.shape
-        let key = key ?? globalState.next()
+        let key = resolve(key: key)
         var result = mlx_array_new()
 
         mlx_random_truncated_normal(
@@ -517,11 +602,11 @@ public enum MLXRandom {
     /// ```
     public static func truncatedNormal(
         low: ScalarOrArray, high: ScalarOrArray, _ shape: [Int]? = nil, dtype: DType = .float32,
-        key: MLXArray? = nil, stream: StreamOrDevice = .default
+        key: RandomStateOrKey? = nil, stream: StreamOrDevice = .default
     ) -> MLXArray {
         let (low, high) = toArrays(low, high)
         let shape = shape ?? low.shape
-        let key = key ?? globalState.next()
+        let key = resolve(key: key)
         var result = mlx_array_new()
 
         mlx_random_truncated_normal(
@@ -547,10 +632,10 @@ public enum MLXRandom {
     /// let array = MLXRandom.gumbel([10, 5], key: key)
     /// ```
     public static func gumbel<T>(
-        _ shape: [Int] = [], type: T.Type = Float.self, key: MLXArray? = nil,
+        _ shape: [Int] = [], type: T.Type = Float.self, key: RandomStateOrKey? = nil,
         stream: StreamOrDevice = .default
     ) -> MLXArray where T: HasDType, T: BinaryFloatingPoint {
-        let key = key ?? globalState.next()
+        let key = resolve(key: key)
         var result = mlx_array_new()
 
         mlx_random_gumbel(
@@ -574,10 +659,10 @@ public enum MLXRandom {
     /// let array = MLXRandom.gumbel([10, 5], key: key)
     /// ```
     public static func gumbel(
-        _ shape: [Int] = [], dtype: DType = .float32, key: MLXArray? = nil,
+        _ shape: [Int] = [], dtype: DType = .float32, key: RandomStateOrKey? = nil,
         stream: StreamOrDevice = .default
     ) -> MLXArray {
-        let key = key ?? globalState.next()
+        let key = resolve(key: key)
         var result = mlx_array_new()
 
         mlx_random_gumbel(&result, shape.asInt32, shape.count, dtype.cmlxDtype, key.ctx, stream.ctx)
@@ -604,10 +689,10 @@ public enum MLXRandom {
     /// - Parameters:
     ///     - logits: The *unnormalized* categorical distribution(s).
     public static func categorical(
-        _ logits: MLXArray, axis: Int = -1, shape: [Int]? = nil, key: MLXArray? = nil,
+        _ logits: MLXArray, axis: Int = -1, shape: [Int]? = nil, key: RandomStateOrKey? = nil,
         stream: StreamOrDevice = .default
     ) -> MLXArray {
-        let key = key ?? globalState.next()
+        let key = resolve(key: key)
         if let shape {
             var result = mlx_array_new()
 
@@ -640,10 +725,10 @@ public enum MLXRandom {
     /// - Parameters:
     ///     - logits: The *unnormalized* categorical distribution(s).
     public static func categorical(
-        _ logits: MLXArray, axis: Int = -1, count: Int, key: MLXArray? = nil,
+        _ logits: MLXArray, axis: Int = -1, count: Int, key: RandomStateOrKey? = nil,
         stream: StreamOrDevice = .default
     ) -> MLXArray {
-        let key = key ?? globalState.next()
+        let key = resolve(key: key)
         var result = mlx_array_new()
 
         mlx_random_categorical_num_samples(
@@ -661,9 +746,9 @@ public enum MLXRandom {
     ///   - scale: scale "b" of the distribution
     public static func laplace(
         _ shape: [Int] = [], dtype: DType = .float32, loc: Float = 0, scale: Float = 1,
-        key: MLXArray? = nil, stream: StreamOrDevice = .default
+        key: RandomStateOrKey? = nil, stream: StreamOrDevice = .default
     ) -> MLXArray {
-        let key = key ?? globalState.next()
+        let key = resolve(key: key)
         var result = mlx_array_new()
 
         mlx_random_laplace(
@@ -724,7 +809,7 @@ public func split(key: MLXArray, stream: StreamOrDevice = .default) -> (MLXArray
 /// let array = MLXRandom.uniform(0.5 ..< 1, [50], key: key)
 /// ```
 public func uniform<R: HasDType, T>(
-    _ range: Range<R>, _ shape: [Int] = [], type: T.Type = Float.self, key: MLXArray? = nil,
+    _ range: Range<R>, _ shape: [Int] = [], type: T.Type = Float.self, key: RandomStateOrKey? = nil,
     stream: StreamOrDevice = .default
 ) -> MLXArray where T: HasDType, T: BinaryFloatingPoint {
     return MLXRandom.uniform(range, shape, type: type, key: key, stream: stream)
@@ -740,7 +825,7 @@ public func uniform<R: HasDType, T>(
 /// ```
 public func uniform<T>(
     _ range: Range<Float> = 0 ..< 1, _ shape: [Int] = [], type: T.Type = Float.self,
-    key: MLXArray? = nil, stream: StreamOrDevice = .default
+    key: RandomStateOrKey? = nil, stream: StreamOrDevice = .default
 ) -> MLXArray where T: HasDType, T: BinaryFloatingPoint {
     return MLXRandom.uniform(range, shape, type: type, key: key, stream: stream)
 }
@@ -761,7 +846,7 @@ public func uniform<T>(
 /// ```
 public func uniform<T>(
     low: ScalarOrArray, high: ScalarOrArray, _ shape: [Int]? = nil, type: T.Type = Float.self,
-    key: MLXArray? = nil, stream: StreamOrDevice = .default
+    key: RandomStateOrKey? = nil, stream: StreamOrDevice = .default
 ) -> MLXArray where T: HasDType, T: BinaryFloatingPoint {
     return MLXRandom.uniform(low: low, high: high, shape, type: type, key: key, stream: stream)
 }
@@ -782,7 +867,7 @@ public func uniform<T>(
 /// ```
 public func uniform(
     low: ScalarOrArray, high: ScalarOrArray, _ shape: [Int]? = nil, dtype: DType = .float32,
-    key: MLXArray? = nil, stream: StreamOrDevice = .default
+    key: RandomStateOrKey? = nil, stream: StreamOrDevice = .default
 ) -> MLXArray {
     return MLXRandom.uniform(low: low, high: high, shape, dtype: dtype, key: key, stream: stream)
 }
@@ -810,7 +895,7 @@ public func uniform(
 ///   - key: PRNG key
 public func normal<T>(
     _ shape: [Int] = [], type: T.Type = Float.self, loc: Float = 0, scale: Float = 1,
-    key: MLXArray? = nil,
+    key: RandomStateOrKey? = nil,
     stream: StreamOrDevice = .default
 ) -> MLXArray where T: HasDType, T: BinaryFloatingPoint {
     return MLXRandom.normal(shape, type: type, loc: loc, scale: scale, key: key, stream: stream)
@@ -839,7 +924,7 @@ public func normal<T>(
 ///   - key: PRNG key
 public func normal(
     _ shape: [Int] = [], dtype: DType = .float32, loc: Float = 0, scale: Float = 1,
-    key: MLXArray? = nil,
+    key: RandomStateOrKey? = nil,
     stream: StreamOrDevice = .default
 ) -> MLXArray {
     return MLXRandom.normal(shape, dtype: dtype, loc: loc, scale: scale, key: key, stream: stream)
@@ -863,7 +948,7 @@ public func normal(
 ///   - key: PRNG key
 public func multivariateNormal(
     mean: MLXArray, covariance: MLXArray, shape: [Int] = [], dtype: DType = .float32,
-    key: MLXArray? = nil, stream: StreamOrDevice = .default
+    key: RandomStateOrKey? = nil, stream: StreamOrDevice = .default
 ) -> MLXArray {
     return MLXRandom.multivariateNormal(
         mean: mean, covariance: covariance, shape: shape, dtype: dtype, key: key, stream: stream)
@@ -885,7 +970,8 @@ public func multivariateNormal(
 /// let array = MLXRandom.randInt(Int32(0) ..< 100, [50], key: key)
 /// ```
 public func randInt<T>(
-    _ range: Range<T>, _ shape: [Int] = [], key: MLXArray? = nil, stream: StreamOrDevice = .default
+    _ range: Range<T>, _ shape: [Int] = [], key: RandomStateOrKey? = nil,
+    stream: StreamOrDevice = .default
 ) -> MLXArray where T: HasDType, T: BinaryInteger {
     return MLXRandom.randInt(range, shape, key: key, stream: stream)
 }
@@ -904,7 +990,7 @@ public func randInt<T>(
 /// let array = MLXRandom.randInt(low: [0, 10], high: [10, 100], key: key)
 /// ```
 public func randInt(
-    low: ScalarOrArray, high: ScalarOrArray, _ shape: [Int]? = nil, key: MLXArray? = nil,
+    low: ScalarOrArray, high: ScalarOrArray, _ shape: [Int]? = nil, key: RandomStateOrKey? = nil,
     stream: StreamOrDevice = .default
 ) -> MLXArray {
     return MLXRandom.randInt(low: low, high: high, shape, key: key, stream: stream)
@@ -926,7 +1012,7 @@ public func randInt(
 /// ```
 public func randInt<T>(
     low: ScalarOrArray, high: ScalarOrArray, _ shape: [Int]? = nil, type: T.Type,
-    key: MLXArray? = nil, stream: StreamOrDevice = .default
+    key: RandomStateOrKey? = nil, stream: StreamOrDevice = .default
 ) -> MLXArray where T: HasDType, T: BinaryInteger {
     return MLXRandom.randInt(low: low, high: high, shape, type: type, key: key, stream: stream)
 }
@@ -945,7 +1031,9 @@ public func randInt<T>(
 /// // generate an array of shape [50, 2] of random Bool
 /// let array = MLXRandom.bernoulli([50, 2], key: key)
 /// ```
-public func bernoulli(_ shape: [Int] = [], key: MLXArray? = nil, stream: StreamOrDevice = .default)
+public func bernoulli(
+    _ shape: [Int] = [], key: RandomStateOrKey? = nil, stream: StreamOrDevice = .default
+)
     -> MLXArray
 {
     return MLXRandom.bernoulli(shape, key: key, stream: stream)
@@ -970,7 +1058,7 @@ public func bernoulli(_ shape: [Int] = [], key: MLXArray? = nil, stream: StreamO
 /// let array = MLXRandom.bernoulli(MLXArray(convert: [0.1, 0.5, 0.8]), key: key)
 /// ```
 public func bernoulli(
-    _ p: ScalarOrArray, _ shape: [Int]? = nil, key: MLXArray? = nil,
+    _ p: ScalarOrArray, _ shape: [Int]? = nil, key: RandomStateOrKey? = nil,
     stream: StreamOrDevice = .default
 ) -> MLXArray {
     return MLXRandom.bernoulli(p, shape, key: key, stream: stream)
@@ -995,7 +1083,7 @@ public func bernoulli(
 /// ### See also
 /// - [JAX Documentation](https://jax.readthedocs.io/en/latest/_modules/jax/_src/random.html#truncated_normal)
 public func truncatedNormal<R: HasDType, T>(
-    _ range: Range<R>, _ shape: [Int] = [], type: T.Type = Float.self, key: MLXArray? = nil,
+    _ range: Range<R>, _ shape: [Int] = [], type: T.Type = Float.self, key: RandomStateOrKey? = nil,
     stream: StreamOrDevice = .default
 ) -> MLXArray where T: HasDType, T: BinaryFloatingPoint {
     return MLXRandom.truncatedNormal(range, shape, type: type, key: key, stream: stream)
@@ -1010,7 +1098,8 @@ public func truncatedNormal<R: HasDType, T>(
 /// let array = MLXRandom.truncatedNormal(0.5 ..< 1, [50], key: key)
 /// ```
 public func truncatedNormal<T>(
-    _ range: Range<Float>, _ shape: [Int] = [], type: T.Type = Float.self, key: MLXArray? = nil,
+    _ range: Range<Float>, _ shape: [Int] = [], type: T.Type = Float.self,
+    key: RandomStateOrKey? = nil,
     stream: StreamOrDevice = .default
 ) -> MLXArray where T: HasDType, T: BinaryFloatingPoint {
     return MLXRandom.truncatedNormal(range, shape, type: type, key: key, stream: stream)
@@ -1031,7 +1120,7 @@ public func truncatedNormal<T>(
 /// ```
 public func truncatedNormal<T>(
     low: ScalarOrArray, high: ScalarOrArray, _ shape: [Int]? = nil, type: T.Type = Float.self,
-    key: MLXArray? = nil, stream: StreamOrDevice = .default
+    key: RandomStateOrKey? = nil, stream: StreamOrDevice = .default
 ) -> MLXArray where T: HasDType, T: BinaryFloatingPoint {
     return MLXRandom.truncatedNormal(
         low: low, high: high, shape, type: type, key: key, stream: stream)
@@ -1052,7 +1141,7 @@ public func truncatedNormal<T>(
 /// ```
 public func truncatedNormal(
     low: ScalarOrArray, high: ScalarOrArray, _ shape: [Int]? = nil, dtype: DType = .float32,
-    key: MLXArray? = nil, stream: StreamOrDevice = .default
+    key: RandomStateOrKey? = nil, stream: StreamOrDevice = .default
 ) -> MLXArray {
     return MLXRandom.truncatedNormal(
         low: low, high: high, shape, dtype: dtype, key: key, stream: stream)
@@ -1073,7 +1162,7 @@ public func truncatedNormal(
 /// let array = MLXRandom.gumbel([10, 5], key: key)
 /// ```
 public func gumbel<T>(
-    _ shape: [Int] = [], type: T.Type = Float.self, key: MLXArray? = nil,
+    _ shape: [Int] = [], type: T.Type = Float.self, key: RandomStateOrKey? = nil,
     stream: StreamOrDevice = .default
 ) -> MLXArray where T: HasDType, T: BinaryFloatingPoint {
     return MLXRandom.gumbel(shape, type: type, key: key, stream: stream)
@@ -1094,7 +1183,7 @@ public func gumbel<T>(
 /// let array = MLXRandom.gumbel([10, 5], key: key)
 /// ```
 public func gumbel(
-    _ shape: [Int] = [], dtype: DType = .float32, key: MLXArray? = nil,
+    _ shape: [Int] = [], dtype: DType = .float32, key: RandomStateOrKey? = nil,
     stream: StreamOrDevice = .default
 ) -> MLXArray {
     return MLXRandom.gumbel(shape, dtype: dtype, key: key, stream: stream)
@@ -1119,7 +1208,7 @@ public func gumbel(
 /// - Parameters:
 ///     - logits: The *unnormalized* categorical distribution(s).
 public func categorical(
-    _ logits: MLXArray, axis: Int = -1, shape: [Int]? = nil, key: MLXArray? = nil,
+    _ logits: MLXArray, axis: Int = -1, shape: [Int]? = nil, key: RandomStateOrKey? = nil,
     stream: StreamOrDevice = .default
 ) -> MLXArray {
     return MLXRandom.categorical(logits, axis: axis, shape: shape, key: key, stream: stream)
@@ -1142,7 +1231,7 @@ public func categorical(
 /// - Parameters:
 ///     - logits: The *unnormalized* categorical distribution(s).
 public func categorical(
-    _ logits: MLXArray, axis: Int = -1, count: Int, key: MLXArray? = nil,
+    _ logits: MLXArray, axis: Int = -1, count: Int, key: RandomStateOrKey? = nil,
     stream: StreamOrDevice = .default
 ) -> MLXArray {
     return MLXRandom.categorical(logits, axis: axis, count: count, key: key, stream: stream)
@@ -1157,7 +1246,7 @@ public func categorical(
 ///   - scale: scale "b" of the distribution
 public func laplace(
     _ shape: [Int] = [], dtype: DType = .float32, loc: Float = 0, scale: Float = 1,
-    key: MLXArray? = nil, stream: StreamOrDevice = .default
+    key: RandomStateOrKey? = nil, stream: StreamOrDevice = .default
 ) -> MLXArray {
     return MLXRandom.laplace(shape, dtype: dtype, loc: loc, scale: scale, key: key, stream: stream)
 }
diff --git a/Source/MLX/State.swift b/Source/MLX/State.swift
index 460e4f56..63e59aae 100644
--- a/Source/MLX/State.swift
+++ b/Source/MLX/State.swift
@@ -2,27 +2,52 @@
 
 import Foundation
 
+/// Protocol for types that can be used as a provider of random keys, e.g. for ``MLXRandom``.
+public protocol RandomStateOrKey {
+    func asRandomKey() -> MLXArray
+}
+
+extension MLXArray: RandomStateOrKey {
+    public func asRandomKey() -> MLXArray {
+        self
+    }
+}
+
 extension MLXRandom {
-    /// Global random state.
+
+    /// Random state factory.
+    ///
+    ///
     ///
     /// Note: although this type is thread-safe, the MLXArrays that it returns are not -- do not
     /// evaluate these values or expressions that depend on them across multiple threads
     /// simultaneously.
-    public class RandomState: Updatable, Evaluatable, @unchecked (Sendable) {
+    ///
+    /// ### See Also
+    /// - ``globalState``
+    /// - ``withRandomState(_:body:)-6i2p1``
+    public class RandomState: RandomStateOrKey, Updatable, Evaluatable, @unchecked (Sendable) {
         private var state: MLXArray
         private let lock = NSLock()
 
-        init() {
+        /// Initialize the RandomState with a seed based on the current time.
+        public init() {
             let now = mach_approximate_time()
             state = MLXRandom.key(now)
         }
 
+        /// Initialize the RandomState with the given seed value.
+        public init(seed: UInt64) {
+            state = MLXRandom.key(seed)
+        }
+
         public func innerState() -> [MLXArray] {
             lock.withLock {
                 [state]
             }
         }
 
+        /// Split the current state and return a new Key.
         public func next() -> MLXArray {
             lock.withLock {
                 let (a, b) = MLXRandom.split(key: state)
@@ -31,11 +56,16 @@ extension MLXRandom {
             }
         }
 
+        /// Reset the random state.
         public func seed(_ seed: UInt64) {
             lock.withLock {
                 state = MLXRandom.key(seed)
             }
         }
+
+        public func asRandomKey() -> MLXArray {
+            next()
+        }
     }
 
     /// Global random state.
@@ -46,4 +76,33 @@ extension MLXRandom {
     /// - ``seed(_:)``
     public static let globalState = RandomState()
 
+    /// See ``withRandomState(_:body:)`` and ``resolve(key:)``
+    @TaskLocal
+    static fileprivate var taskLocalRandomState: MLXRandom.RandomState?
+
 }  // MLXRandom
+
+/// Resolve the given key to a concrete MLXArray (random key).
+///
+/// This will use the following values (in order until one is found) to resolve the
+/// random key:
+///
+/// - the passed key, either an ``MLXArray`` or ``MLXRandom/RandomState``
+/// - the task-local ``MLXRandom/RandomState``, see ``withRandomState(_:body:)-6i2p1``
+/// - the global RandomState, ``MLXRandom/globalState``
+public func resolve(key: RandomStateOrKey?) -> MLXArray {
+    key?.asRandomKey() ?? MLXRandom.taskLocalRandomState?.asRandomKey()
+        ?? MLXRandom.globalState.next()
+}
+
+/// Use the given ``MLXRandom/RandomState`` scoped to the current task and body.
+public func withRandomState<R>(_ state: MLXRandom.RandomState, body: () throws -> R) rethrows -> R {
+    try MLXRandom.$taskLocalRandomState.withValue(state, operation: body)
+}
+
+/// Use the given ``MLXRandom/RandomState`` scoped to the current task and body.
+public func withRandomState<R>(_ state: MLXRandom.RandomState, body: () async throws -> R)
+    async rethrows -> R
+{
+    try await MLXRandom.$taskLocalRandomState.withValue(state, operation: body)
+}
diff --git a/Tests/MLXTests/MLXRandomTests.swift b/Tests/MLXTests/MLXRandomTests.swift
index 8613c360..07a3125f 100644
--- a/Tests/MLXTests/MLXRandomTests.swift
+++ b/Tests/MLXTests/MLXRandomTests.swift
@@ -159,4 +159,106 @@ class MLXRandomTests: XCTestCase {
         assertEqual(result, expected)
     }
 
+    func testRandomStateOrKeySame() {
+        // these should all produce the same value since they
+        // all resolve to the same key
+
+        let key = MLXRandom.key(0)
+        let (_, k1) = split(key: key)
+
+        let state = MLXRandom.RandomState(seed: 0)
+        MLXRandom.seed(0)
+
+        // global state
+        let v0 = uniform(0 ..< 1, [5])
+
+        // explicit key
+        let v1 = uniform(0 ..< 1, [5], key: k1)
+
+        // local RandomState
+        let v2 = uniform(0 ..< 1, [5], key: state)
+
+        assertEqual(v0, v1)
+        assertEqual(v1, v2)
+    }
+
+    func testRandomStateOrKeyDifferent() {
+        // these should all produce different values as they
+        // use different keys -- note this is otherwise identical
+        // to testRandomStateOrKeySame
+
+        let key = MLXRandom.key(7)
+        let (_, k1) = split(key: key)
+
+        let state = MLXRandom.RandomState(seed: 11)
+        MLXRandom.seed(31)
+
+        // global state
+        let v0 = uniform(0 ..< 1, [5])
+
+        // explicit key
+        let v1 = uniform(0 ..< 1, [5], key: k1)
+
+        // local RandomState
+        let v2 = uniform(0 ..< 1, [5], key: state)
+
+        assertNotEqual(v0, v1)
+        assertNotEqual(v1, v2)
+        assertNotEqual(v0, v2)
+    }
+
+    func testRandomThreadsSame() async {
+        // several threads using task local random state with a constant
+        // seed will produce the same value
+        await withTaskGroup(of: Float.self) { group in
+            for _ in 0 ..< 10 {
+                group.addTask {
+                    let state = MLXRandom.RandomState(seed: 23)
+                    return withRandomState(state) {
+                        var t: Float = 0.0
+                        for _ in 0 ..< 100 {
+                            t += uniform(0 ..< 1, [10, 10]).sum().item(Float.self)
+                        }
+                        return t
+                    }
+                }
+            }
+
+            var result = [Float]()
+            for await v in group {
+                result.append(v)
+            }
+
+            let unique = Set(result)
+            XCTAssertEqual(unique.count, 1, "Different values: \(result)")
+        }
+    }
+
+    func testRandomThreadsDifferent() async {
+        // several threads using task local random state with different
+        // seeds will produce different values
+        await withTaskGroup(of: Float.self) { group in
+            for i in 0 ..< 10 {
+                group.addTask {
+                    let state = MLXRandom.RandomState(seed: UInt64(i))
+                    return withRandomState(state) {
+                        var t: Float = 0.0
+                        for _ in 0 ..< 100 {
+                            t += uniform(0 ..< 1, [10, 10]).sum().item(Float.self)
+                        }
+                        return t
+                    }
+                }
+            }
+
+            var result = [Float]()
+            for await v in group {
+                result.append(v)
+            }
+
+            let unique = Set(result)
+            XCTAssertEqual(unique.count, 10, "Same values: \(result)")
+        }
+    }
+
 }
diff --git a/Tests/MLXTests/Utils.swift b/Tests/MLXTests/Utils.swift
index 73e58e23..520306bd 100644
--- a/Tests/MLXTests/Utils.swift
+++ b/Tests/MLXTests/Utils.swift
@@ -24,6 +24,16 @@ func assertEqual(
     }
 }
 
+func assertNotEqual(
+    _ array1: MLXArray, _ array2: MLXArray, rtol: Double = 1e-5, atol: Double = 1e-8,
+    file: StaticString = #filePath, line: UInt = #line
+) {
+    XCTAssertEqual(array1.shape, array2.shape, "shapes differ: \(array1.shape) != \(array2.shape)")
+    XCTAssertFalse(
+        array1.allClose(array2, rtol: rtol, atol: atol).item(Bool.self),
+        "contents same:\n\(array1)\n\(array2)")
+}
+
 func setDefaultDevice() {
     MLX.Device.setDefault(device: .gpu)
 }


Reviews:
user: login: davidkoski
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-08T19:47:57Z
user: login: davidkoski
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-08T19:48:36Z
user: login: davidkoski
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-08T19:49:10Z
user: login: davidkoski
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-08T19:51:25Z
user: login: davidkoski
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-08T19:52:19Z
user: login: awni
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-09T13:53:27Z
user: login: awni
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-18T13:14:07Z
user: login: awni
body: I'm ok merging this. I think my hesitation is that the `randomKeyOrState` in the random API is causes some divergence from the rest of the MLX APIs (Python / C++). Other than that I don't have a problem with it and I can see how it'd be useful.

@davidkoski I leave it up to you to keep or not depending on how useful you think it is.
state: APPROVED
author_association: MEMBER
submitted_at: 2025-04-18T13:16:40Z

Review comments:
diff_hunk: @@ -2,21 +2,36 @@
 
 import Foundation
 
+public protocol RandomStateOrKey {
path: Source/MLX/State.swift
user: login: davidkoski
body: - [x] various need documentation
created_at: 2025-04-08T19:47:57Z
author_association: COLLABORATOR
diff_hunk: @@ -67,12 +67,13 @@ public enum MLXRandom {
     /// let array = MLXRandom.uniform(0.5 ..< 1, [50], key: key)
     /// ```
     public static func uniform<R: HasDType, T>(
-        _ range: Range<R>, _ shape: [Int] = [], type: T.Type = Float.self, key: MLXArray? = nil,
+        _ range: Range<R>, _ shape: [Int] = [], type: T.Type = Float.self,
+        key: RandomStateOrKey? = nil,
path: Source/MLX/Random.swift
user: login: davidkoski
body: allow passing a Key (MLXArray) or RandomState (vendor of keys)
created_at: 2025-04-08T19:48:36Z
author_association: COLLABORATOR
diff_hunk: @@ -2,21 +2,36 @@
 
 import Foundation
 
+public protocol RandomStateOrKey {
+    func asRandomKey() -> MLXArray
+}
+
+extension MLXArray: RandomStateOrKey {
+    public func asRandomKey() -> MLXArray {
+        self
+    }
+}
+
 extension MLXRandom {
+
     /// Global random state.
     ///
     /// Note: although this type is thread-safe, the MLXArrays that it returns are not -- do not
     /// evaluate these values or expressions that depend on them across multiple threads
     /// simultaneously.
-    public class RandomState: Updatable, Evaluatable, @unchecked (Sendable) {
+    public class RandomState: RandomStateOrKey, Updatable, Evaluatable, @unchecked (Sendable) {
         private var state: MLXArray
         private let lock = NSLock()
 
-        init() {
+        public init() {
path: Source/MLX/State.swift
user: login: davidkoski
body: These inits should be public so that callers can create these as needed.
created_at: 2025-04-08T19:49:10Z
author_association: COLLABORATOR
diff_hunk: @@ -47,3 +66,20 @@ extension MLXRandom {
     public static let globalState = RandomState()
 
 }  // MLXRandom
+
+@TaskLocal
+private var taskLocalRandomState: MLXRandom.RandomState?
+
+public func resolve(key: RandomStateOrKey?) -> MLXArray {
+    key?.asRandomKey() ?? taskLocalRandomState?.asRandomKey() ?? MLXRandom.globalState.next()
path: Source/MLX/State.swift
user: login: davidkoski
body: This lets callers of random functions use:

- a passed in Key or RandomState
- a task (roughly thread) local RandomState, see `withRandomState` in tests
- the global RandomState

Previously this was either a passed in MLXArray OR the global RandomState, which made it hard to control if the use of MLXRandom was buried deeply in a model or sampling.
created_at: 2025-04-08T19:51:25Z
author_association: COLLABORATOR
diff_hunk: @@ -159,4 +159,106 @@ class MLXRandomTests: XCTestCase {
         assertEqual(result, expected)
     }
 
+    func testRandomStateOrKeySame() {
+        // these should all produce the same value since they
+        // all resolve to the same key
+
+        let key = MLXRandom.key(0)
+        let (_, k1) = split(key: key)
+
+        let state = MLXRandom.RandomState(seed: 0)
+        MLXRandom.seed(0)
+
+        // global state
+        let v0 = uniform(0 ..< 1, [5])
+
+        // explicit key
+        let v1 = uniform(0 ..< 1, [5], key: k1)
+
+        // local RandomState
+        let v2 = uniform(0 ..< 1, [5], key: state)
+
+        assertEqual(v0, v1)
+        assertEqual(v1, v2)
+    }
+
+    func testRandomStateOrKeyDifferent() {
+        // these should all produce different values as they
+        // use different keys -- note this is otherwise identical
+        // to testRandomStateOrKeySame
+
+        let key = MLXRandom.key(7)
+        let (_, k1) = split(key: key)
+
+        let state = MLXRandom.RandomState(seed: 11)
+        MLXRandom.seed(31)
+
+        // global state
+        let v0 = uniform(0 ..< 1, [5])
+
+        // explicit key
+        let v1 = uniform(0 ..< 1, [5], key: k1)
+
+        // local RandomState
+        let v2 = uniform(0 ..< 1, [5], key: state)
+
+        assertNotEqual(v0, v1)
+        assertNotEqual(v1, v2)
+        assertNotEqual(v0, v2)
+    }
+
+    func testRandomThreadsSame() async {
+        // several threads using task local random state with a constant
+        // seed will produce the same value
+        await withTaskGroup { group in
+            for _ in 0 ..< 10 {
+                group.addTask {
+                    let state = MLXRandom.RandomState(seed: 23)
+                    return withRandomState(state) {
path: Tests/MLXTests/MLXRandomTests.swift
user: login: davidkoski
body: This will use per-task (thread) RandomState -- this would crash if done before this change.
created_at: 2025-04-08T19:52:19Z
author_association: COLLABORATOR
diff_hunk: @@ -3,6 +3,70 @@
 import Cmlx
 import Foundation
 
+/// Collection of functions related to random number generation.
+///
+/// Following [JAX’s PRNG design](https://jax.readthedocs.io/en/latest/jep/263-prng.html)
+/// we use a splittable version of Threefry, which is a counter-based PRNG.
+///
+/// Random sampling functions in MLX use an implicit global PRNG state by default.
+/// However, all functions take an optional key keyword argument for when more fine-grained
+/// control or explicit state management is needed.  Callers can also arrange for `Task` local
+/// random state -- useful in multithreaded situations.
+///
+/// For example, you can generate random numbers with:
+///
+/// ```swift
+/// for _ in 0 ..< 3 {
+///   print(MLXRandom.uniform())
+/// }
+/// ```
+///
+/// which will print a sequence of unique pseudo random numbers. Alternatively you can explicitly set the key:
+///
+/// ```swift
+/// let key = MLXRandom.key(0)
+/// for _ in 0 ..< 3 {
+///   print(MLXRandom.uniform(key: key))
+/// }
+/// ```
+///
+/// which will yield the same pseudo random number at each iteration as the key doesn't change.
+///
+/// You can also use a ``RandomState`` to generate different random numbers but the same sequence
+/// each time:
+///
+/// ```swift
+/// let state = RandomState(seed: 0)
+/// for _ in 0 ..< 3 {
+///   print(MLXRandom.uniform(key: state))
+/// }
path: Source/MLX/Random.swift
user: login: awni
body: I think it would be helpful in the exposition here to add an example inbetween this one and the previous one that explains how to split the array random key. So concretely:

1. Example with array key -> same numbers
2. Example with array key + split -> different numbers
3. Same thing as 2 but nicer syntax with random state
created_at: 2025-04-09T13:53:27Z
author_association: MEMBER
diff_hunk: @@ -46,4 +76,33 @@ extension MLXRandom {
     /// - ``seed(_:)``
     public static let globalState = RandomState()
 
+    /// See ``withRandomState(_:body:)`` and ``resolve(key:)``
+    @TaskLocal
+    static fileprivate var taskLocalRandomState: MLXRandom.RandomState?
+
 }  // MLXRandom
+
+/// Resolve the given key to a concrete MLXArray (random key).
+///
+/// This will use the following values (in order until one is found) to resolve the
+/// random key:
+///
+/// - the passed key, either an ``MLXArray`` or ``MLXRandom/RandomState``
+/// - the task-local ``MLXRandom/RandomState``, see ``withRandomState(_:body:)-6i2p1``
+/// - the global RandomState, ``MLXRandom/globalState``
+public func resolve(key: RandomStateOrKey?) -> MLXArray {
+    key?.asRandomKey() ?? MLXRandom.taskLocalRandomState?.asRandomKey()
+        ?? MLXRandom.globalState.next()
+}
path: Source/MLX/State.swift
user: login: awni
body: Why did you make that public? Shouldn't it not be an implementation only function?
created_at: 2025-04-18T13:14:07Z
author_association: MEMBER

PR comments:
user: login: awni
created_at: 2025-04-09T13:59:28Z
body: @davidkoski right now I see that the random state in Swift is already "thread safe" as you guard the state with a lock. Is that correct?
author_association: MEMBER
user: login: awni
created_at: 2025-04-09T14:03:55Z
body: > allow a Key or a RandomState to be used with Random functions

This one I'm not sure about. Can you say more about the motivation for adding that?

I kind of feel like for people that want control over the individual calls.. they can use the splittable key. It's slightly more verbose but it's also an edge case.. And then for people that don't need control over individual calls, they can create the state and use it in a local context and/or seed the global state.
author_association: MEMBER
user: login: davidkoski
created_at: 2025-04-09T14:59:23Z
body: > @davidkoski right now I see that the random state in Swift is already "thread safe" as you guard the state with a lock. Is that correct?

Sort of -- although the container itself is thread safe (so I can mark it Sendable), the MLXArrays that it vends are not thread safe.  It is kind of a useless distinction but one I need to be able to send these between threads (the `@TaskLocal` requires it for instance).
author_association: COLLABORATOR
user: login: davidkoski
created_at: 2025-04-09T15:02:47Z
body: > > allow a Key or a RandomState to be used with Random functions
> 
> This one I'm not sure about. Can you say more about the motivation for adding that?
> 
> I kind of feel like for people that want control over the individual calls.. they can use the splittable key. It's slightly more verbose but it's also an edge case.. And then for people that don't need control over individual calls, they can create the state and use it in a local context and/or seed the global state.

If you wanted to have a private random state for something, e.g. a categorical sampler in a LLM inference pipeline, you would have to use the key split and manage state that way.  This is just a convenience if you want to have a random state without the need to manually manage the splitting.

I don't think it is critical for this change -- the `@TaskLocal` to inject a scoped "global" random state is the more important part, but I do think that using `RandomState` is more convenient than manually managing the splits (yet equivalent).

I think it is exactly as you said in point 3:

> - Example with array key -> same numbers
> - Example with array key + split -> different numbers
> - Same thing as 2 but nicer syntax with random state

This:

```swift
var key = MLXRandom.key(0)
for _ in 0 ..< 3 {
  let (a, b) = MLXRandom.split(key: key)
  // use b to generate a different value each time
  print(MLXRandom.uniform(key: b))
  // new random state is a
  key = a
}
```

vs

```swift
let state = RandomState(seed: 0)
for _ in 0 ..< 3 {
  print(MLXRandom.uniform(key: state))
}
```

They are exactly the same operations.
author_association: COLLABORATOR
user: login: davidkoski
created_at: 2025-04-18T16:53:59Z
body: OK, let's go with what we have here then -- I think passing RandomState is a nice alternative to using implicit TaskLocals, for those who desire explicit state instead.  We can see if this appeals to anyone and it would be easy enough to add on the python/c++ side if there is uptake.

Given current use of Random I suspect most people will just use TaskLocal.  C++ might need that approach, though python probably does not.
author_association: COLLABORATOR