Skip to content

Commit

Permalink
This solves the Numba issue; the Awkward one will need an Awkward fix…
Browse files Browse the repository at this point in the history
… (update ak.nan_to_num).
  • Loading branch information
jpivarski committed Feb 18, 2022
1 parent 33f2418 commit 12ae4d6
Showing 1 changed file with 25 additions and 10 deletions.
35 changes: 25 additions & 10 deletions src/vector/_backends/numba_object.py
Original file line number Diff line number Diff line change
Expand Up @@ -52,16 +52,31 @@
def nan_to_num(x, copy=True, nan=0.0, posinf=None, neginf=None):
if isinstance(x, numba.types.Array):

def nan_to_num_impl(x, copy=True, nan=0.0, posinf=None, neginf=None):
out = numpy.copy(x).reshape(-1) if copy else x.reshape(-1)
for i in range(len(out)):
if numpy.isnan(out[i]):
out[i] = nan
if posinf is not None and numpy.isinf(out[i]) and out[i] > 0:
out[i] = posinf
if neginf is not None and numpy.isinf(out[i]) and out[i] < 0:
out[i] = neginf
return out.reshape(x.shape)
if isinstance(nan, numba.types.Array):

def nan_to_num_impl(x, copy=True, nan=0.0, posinf=None, neginf=None):
out = numpy.copy(x).reshape(-1) if copy else x.reshape(-1)
for i in range(len(out)):
if numpy.isnan(out[i]):
out[i] = nan[i]
if posinf is not None and numpy.isinf(out[i]) and out[i] > 0:
out[i] = posinf
if neginf is not None and numpy.isinf(out[i]) and out[i] < 0:
out[i] = neginf
return out.reshape(x.shape)

else:

def nan_to_num_impl(x, copy=True, nan=0.0, posinf=None, neginf=None):
out = numpy.copy(x).reshape(-1) if copy else x.reshape(-1)
for i in range(len(out)):
if numpy.isnan(out[i]):
out[i] = nan
if posinf is not None and numpy.isinf(out[i]) and out[i] > 0:
out[i] = posinf
if neginf is not None and numpy.isinf(out[i]) and out[i] < 0:
out[i] = neginf
return out.reshape(x.shape)

else:

Expand Down

0 comments on commit 12ae4d6

Please sign in to comment.