tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 11:56:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

clean up xpow (#14295)

Browse Source 
mostly for `ret * (base < 0).where(adj, ret.const_like(1))` -> `(base < 0).where(neg_base, ret)`, since it's good for NAN neg_base but not generic
This commit is contained in:
chenyu
2026-01-23 10:19:47 -05:00
committed by GitHub
parent b913c910c5
commit d5a3b02a9c
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/uop/decompositions.py
+4 -4
View File
@@ -258,12 +258,12 @@ def xlog2(d:UOp) -> UOp:
def xpow(base:UOp, exponent:UOp) -> UOp:
# start with b ** e = exp2(e * log2(b))
ret = (base < 0).where(-base, base).log2().mul(exponent).exp2()
# negative base adjustment: nan for non-integer exponent and -1 for odd exponent
# negative base: nan for non-integer exponent, negate for odd integer exponent
non_int = exponent != exponent.cast(dtypes.int32).cast(exponent.dtype)
adj = non_int.where(ret.const_like(math.nan),
(exponent < 0).where(-exponent, exponent).cast(dtypes.int32).mod(2).cast(dtypes.bool).where(ret.const_like(-1), ret.const_like(1)))
is_odd = (exponent < 0).where(-exponent, exponent).cast(dtypes.int32).mod(2).cast(dtypes.bool)
neg_base = non_int.where(ret.const_like(math.nan), is_odd.where(-ret, ret))
# fix 0 ** 0 = 1
return (base.eq(0) & exponent.eq(0)).where(ret.const_like(1), ret * (base < 0).where(adj, ret.const_like(1)))
return (base.eq(0) & exponent.eq(0)).where(ret.const_like(1), (base < 0).where(neg_base, ret))
# *** integer division ***