Repository: tinygrad/tinygrad

Issue:
Negative tensors taken to power 3 or greater return nan
- Issue: Negative tensors taken to power 3 or greater return nan. Does not match with PyTorch output
- Problem: In tensor.py -> def_pow, function tries to take log of a negative number. Log of a negative number = nan
-Might have been caused by the recent merge "Remove POW llop and add SQRT llop". 

Test code compared with PyTorch:
```
print(torch.tensor(-2).pow(3))
print(Tensor(-2).pow(3).numpy())
-> tensor(-8)
-> nan
```

"def pow" code: https://github.com/geohot/tinygrad/pull/1104/files
tensor.py -> def_pow code:
```
def pow(self, x:Union[Tensor, float], reverse=False) -> Tensor:
    if x.__class__ is not Tensor and not reverse:
      # simple pow identities
      if x < 0: return (1.0/self).pow(-x)
      if x == 2.0: return self*self
      if x == 1.0: return self
      if x == 0.5: return self.sqrt()
    return self.log().mul(x).exp() if not reverse or isinstance(x, Tensor) else self.mul(log(x)).exp()
```


Discussion:
Related: https://github.com/geohot/tinygrad/pull/959/files
Wrote a quick working implementation (I only tested on a few combinations of numbers tho) if someone wants to take it and clean it up and add tests for it, should be a pretty good first pr:
```py
def pow(self, x:Union[Tensor, float], reverse=False) -> Tensor:
  if x.__class__ is not Tensor and not reverse:
    # simple pow identities
    if x < 0: return (1.0/self).pow(-x)
    if x == 2.0: return self*self
    if x == 1.0: return self
    if x == 0.5: return self.sqrt()
  ar = self.abs().log().mul(x).exp() if not reverse or isinstance(x, Tensor) else self.mul(log(abs(x))).exp()
  # correct sign of negative numbers raised to a power (cos has a period of 2pi so we use it here to get the oddness of the power)
  sign = (x * pi).cos() if isinstance(x, Tensor) else cos(x * pi) if not reverse else (self * pi).cos()
  # we only need to correct the sign if the base is negative
  base_sign = ((self.sign() if not reverse else x.sign() if isinstance(x, Tensor) else copysign(1, x)) - 1) / -2
  return ar.mul(sign * base_sign + (1 - base_sign))
```

This is very loosely based off of #959 and the rocm implementation here: https://github.com/RadeonOpenCompute/ROCm-Device-Libs/blob/amd-stg-open/ocml/src/powF_base.h

Resolving patch:
diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py
--- a/tinygrad/tensor.py
+++ b/tinygrad/tensor.py
@@ -7,7 +7,7 @@
 import numpy as np
 from typing import List, Tuple, Callable, Optional, ClassVar, Type, Union, Sequence, cast
 from tinygrad.helpers import ImageDType, argfix, make_pair, getenv, IMAGE, DEBUG, flatten, DType, dtypes
-from math import ceil, pi, prod, sqrt, log
+from math import ceil, pi, prod, sqrt, log, cos, copysign
 from tinygrad.lazy import Device, LazyBuffer
 from tinygrad.ops import LoadOps
 
@@ -561,7 +561,12 @@ def pow(self, x:Union[Tensor, float], reverse=False) -> Tensor:
       if x == 2.0: return self*self
       if x == 1.0: return self
       if x == 0.5: return self.sqrt()
-    return self.log().mul(x).exp() if not reverse or isinstance(x, Tensor) else self.mul(log(x)).exp()
+    ar = self.abs().log().mul(x).exp() if not reverse or isinstance(x, Tensor) else self.mul(log(abs(x))).exp()
+    # correct sign of negative numbers raised to a power (cos has a period of 2pi so we use it here to get the oddness of the power)
+    sign = (x * pi).cos() if isinstance(x, Tensor) else cos(x * pi) if not reverse else (self * pi).cos()
+    # we only need to correct the sign if the base is negative
+    base_sign = ((self.sign() if not reverse else x.sign() if isinstance(x, Tensor) else copysign(1, x)) - 1) / -2
+    return ar.mul(sign * base_sign + (1 - base_sign))
   def matmul(self, x:Tensor, reverse=False) -> Tensor: return x.dot(self) if reverse else self.dot(x)
 
   def maximum(self, x:Union[Tensor, float]) -> Tensor: return self._broadcasted(mlops.Maximum, x)


Test patch:
diff --git a/test/test_ops.py b/test/test_ops.py
--- a/test/test_ops.py
+++ b/test/test_ops.py
@@ -218,6 +218,9 @@ def test_pow(self):
     helper_test_op([(45,65), (45,65)], lambda x,y: x**y, Tensor.pow, a=0)
     helper_test_op([()], lambda x: x**2, lambda x: Tensor.pow(x,2), a=0)
     helper_test_op([()], lambda x: x**-2, lambda x: Tensor.pow(x,-2), a=0)
+    # Regression tests for https://github.com/tinygrad/tinygrad/issues/1151
+    helper_test_op([(45,65)], lambda x: x**3, lambda x: Tensor.pow(x,3), a=-10)
+    helper_test_op([()], lambda x: x**3, lambda x: Tensor.pow(x,3), a=-10)
   def test_pow_const(self):
     helper_test_op([(45,65)], lambda x: x**1.0, lambda x: x**1.0)
     helper_test_op([(45,65)], lambda x: x**-1.0, lambda x: x**-1.0)