From 0904df84e296c10fe0ccb37a3b689e6b96d87e41 Mon Sep 17 00:00:00 2001 From: AUTOMATIC1111 <16777216c@gmail.com> Date: Thu, 3 Aug 2023 07:53:03 +0300 Subject: [PATCH] minor performance improvements for philox --- modules/rng_philox.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/modules/rng_philox.py b/modules/rng_philox.py index b5c024836..5532cf9dd 100644 --- a/modules/rng_philox.py +++ b/modules/rng_philox.py @@ -26,7 +26,7 @@ two_pow32_inv_2pi = np.array([2.3283064e-10 * 6.2831855], dtype=np.float32) def uint32(x): """Converts (N,) np.uint64 array into (2, N) np.unit32 array.""" - return np.moveaxis(x.view(np.uint32).reshape(-1, 2), 0, 1) + return x.view(np.uint32).reshape(-1, 2).transpose(1, 0) def philox4_round(counter, key): @@ -65,8 +65,8 @@ def philox4_32(counter, key, rounds=10): def box_muller(x, y): """Returns just the first out of two numbers generated by Box–Muller transform algorithm.""" - u = x.astype(np.float32) * two_pow32_inv + two_pow32_inv / 2 - v = y.astype(np.float32) * two_pow32_inv_2pi + two_pow32_inv_2pi / 2 + u = x * two_pow32_inv + two_pow32_inv / 2 + v = y * two_pow32_inv_2pi + two_pow32_inv_2pi / 2 s = np.sqrt(-2.0 * np.log(u)) @@ -93,7 +93,9 @@ class Generator: counter[2] = np.arange(n, dtype=np.uint32) # up to 2^32 numbers can be generated - if you want more you'd need to spill into counter[3] self.offset += 1 - key = uint32(np.array([[self.seed] * n], dtype=np.uint64)) + key = np.empty(n, dtype=np.uint64) + key.fill(self.seed) + key = uint32(key) g = philox4_32(counter, key)