mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #25655 from carlosgmartin:simplify_random_orthogonal

Browse Source 
PiperOrigin-RevId: 722770603
This commit is contained in:
jax authors 2025-02-03 13:12:09 -08:00
commit 2c10a65b73
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
jax/_src/random.py
View File

@ -2070,8 +2070,8 @@ def orthogonal(
n = core.concrete_or_error(index, n, "The error occurred in jax.random.orthogonal()")
z = normal(key, (*shape, n, n), dtype)
q, r = jnp.linalg.qr(z)
d = jnp.diagonal(r, 0, -2, -1)
return lax.mul(q, lax.expand_dims(lax.div(d, abs(d).astype(d.dtype)), [-2]))
d = jnp.linalg.diagonal(r)
return q * jnp.expand_dims(jnp.sign(d), -2)
def generalized_normal(
key: ArrayLike,