Skip to content
This repository has been archived by the owner on Dec 6, 2024. It is now read-only.

Commit

Permalink
Refine dtype logic for Power node
Browse files Browse the repository at this point in the history
We know `0 ** x` and `x ** 0`, but we can't assume that `x` is an instance of `Constant` with a defined dtype.
  • Loading branch information
angus-g authored Nov 7, 2024
1 parent 1944432 commit 50eb4a3
Showing 1 changed file with 8 additions and 2 deletions.
10 changes: 8 additions & 2 deletions gem/gem.py
Original file line number Diff line number Diff line change
Expand Up @@ -415,12 +415,18 @@ def __new__(cls, base, exponent):

# Constant folding
if isinstance(base, Zero):
dtype = numpy.result_type(base.dtype, exponent.dtype)
if isinstance(exponent, Constant):
dtype = numpy.result_type(base.dtype, exponent.dtype)
else:
dtype = base.dtype
if isinstance(exponent, Zero):
raise ValueError("cannot solve 0^0")
return Zero(dtype=dtype)
elif isinstance(exponent, Zero):
dtype = numpy.result_type(base.dtype, exponent.dtype)
if isinstance(base, Constant):
dtype = numpy.result_type(base.dtype, exponent.dtype)
else:
dtype = exponent.dtype
return Literal(1, dtype=dtype)
elif isinstance(base, Constant) and isinstance(exponent, Constant):
dtype = numpy.result_type(base.dtype, exponent.dtype)
Expand Down

0 comments on commit 50eb4a3

Please sign in to comment.