Skip to content

Commit

Permalink
Merge branch 'main' into agoose77/docs-emsdk-version
Browse files Browse the repository at this point in the history
  • Loading branch information
@agoose77
agoose77 authored Feb 15, 2023
2 parents 5d0ec0e + 6b49960 commit 9fcae19
Show file tree
Hide file tree
Showing 8 changed files with 492 additions and 136 deletions.
11 changes: 8 additions & 3 deletions src/awkward/_nplikes/array_module.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -99,8 +99,10 @@ def meshgrid(

def array_equal(
self, x1: ArrayLike, x2: ArrayLike, *, equal_nan: bool = False
) -> bool:
return self._module.array_equal(x1, x2, equal_nan=equal_nan)
) -> ArrayLike:
return self._module.asarray(
self._module.array_equal(x1, x2, equal_nan=equal_nan)
)

def searchsorted(
self,
Expand All @@ -118,7 +120,7 @@ def broadcast_arrays(self, *arrays: ArrayLike) -> list[ArrayLike]:
return self._module.broadcast_arrays(*arrays)

def reshape(
self, x: ArrayLike, shape: tuple[int, ...], *, copy: bool | None = None
self, x: ArrayLike, shape: tuple[ShapeItem, ...], *, copy: bool | None = None
) -> ArrayLike:
if copy is False:
raise wrap_error(
Expand Down Expand Up @@ -184,6 +186,9 @@ def regularize_index_for_length(
def nonzero(self, x: ArrayLike) -> tuple[ArrayLike, ...]:
return self._module.nonzero(x)

def where(self, condition: ArrayLike, x1: ArrayLike, x2: ArrayLike) -> ArrayLike:
return self._module.where(condition, x1, x2)

def unique_values(self, x: ArrayLike) -> ArrayLike:
return self._module.unique(
x,
Expand Down
78 changes: 72 additions & 6 deletions src/awkward/_nplikes/numpylike.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,59 @@ def __len__(self) -> int:
def view(self, dtype: dtype) -> Self:
...

# Scalar UFUNCS
@abstractmethod
def __add__(self, other: int | complex | float | Self) -> Self:
...

@abstractmethod
def __sub__(self, other: int | complex | float | Self) -> Self:
...

@abstractmethod
def __mul__(self, other: int | complex | float | Self) -> Self:
...

@abstractmethod
def __truediv__(self, other: int | complex | float | Self) -> Self:
...

@abstractmethod
def __floordiv__(self, other: int | complex | float | Self) -> Self:
...

@abstractmethod
def __gt__(self, other: int | complex | float | Self) -> Self:
...

@abstractmethod
def __lt__(self, other: int | complex | float | Self) -> Self:
...

@abstractmethod
def __ge__(self, other: int | complex | float | Self) -> Self:
...

@abstractmethod
def __le__(self, other: int | complex | float | Self) -> Self:
...

@abstractmethod
def __eq__(self, other: int | complex | float | bool | Self) -> Self:
...

@abstractmethod
def __and__(self, other: int | bool | Self) -> Self:
...

@abstractmethod
def __or__(self, other: int | bool | Self) -> Self:
...

@abstractmethod
def __invert__(self) -> Self:
...


class NumpyMetadata(Singleton):
bool_ = numpy.bool_
Expand Down Expand Up @@ -208,26 +261,35 @@ def frombuffer(

@abstractmethod
def zeros(
self, shape: int | tuple[int, ...], *, dtype: numpy.dtype | None = None
self,
shape: ShapeItem | tuple[ShapeItem, ...],
*,
dtype: numpy.dtype | None = None,
) -> ArrayLike:
...

@abstractmethod
def ones(
self, shape: int | tuple[int, ...], *, dtype: numpy.dtype | None = None
self,
shape: ShapeItem | tuple[ShapeItem, ...],
*,
dtype: numpy.dtype | None = None,
) -> ArrayLike:
...

@abstractmethod
def empty(
self, shape: int | tuple[int, ...], *, dtype: numpy.dtype | None = None
self,
shape: ShapeItem | tuple[ShapeItem, ...],
*,
dtype: numpy.dtype | None = None,
) -> ArrayLike:
...

@abstractmethod
def full(
self,
shape: int | tuple[int, ...],
shape: ShapeItem | tuple[ShapeItem, ...],
fill_value,
*,
dtype: numpy.dtype | None = None,
Expand Down Expand Up @@ -272,7 +334,7 @@ def meshgrid(
@abstractmethod
def array_equal(
self, x1: ArrayLike, x2: ArrayLike, *, equal_nan: bool = False
) -> bool:
) -> ArrayLike:
...

@abstractmethod
Expand Down Expand Up @@ -326,14 +388,18 @@ def regularize_index_for_length(

@abstractmethod
def reshape(
self, x: ArrayLike, shape: tuple[int, ...], *, copy: bool | None = None
self, x: ArrayLike, shape: tuple[ShapeItem, ...], *, copy: bool | None = None
) -> ArrayLike:
...

@abstractmethod
def nonzero(self, x: ArrayLike) -> tuple[ArrayLike, ...]:
...

@abstractmethod
def where(self, condition: ArrayLike, x1: ArrayLike, x2: ArrayLike) -> ArrayLike:
...

@abstractmethod
def unique_values(self, x: ArrayLike) -> ArrayLike:
...
Expand Down
32 changes: 23 additions & 9 deletions src/awkward/_nplikes/typetracer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,14 @@
from awkward._nplikes.numpylike import ArrayLike, IndexType, NumpyLike, NumpyMetadata
from awkward._nplikes.shape import ShapeItem, unknown_length
from awkward._util import NDArrayOperatorsMixin, is_non_string_like_sequence
from awkward.typing import Any, Final, Literal, Self, SupportsIndex, TypeVar
from awkward.typing import (
Any,
Final,
Literal,
Self,
SupportsIndex,
TypeVar,
)

np = NumpyMetadata.instance()

Expand Down Expand Up @@ -764,10 +771,10 @@ def meshgrid(

def array_equal(
self, x1: ArrayLike, x2: ArrayLike, *, equal_nan: bool = False
) -> bool:
) -> TypeTracerArray:
try_touch_data(x1)
try_touch_data(x2)
return False
return TypeTracerArray._new(np.bool_, shape=())

def searchsorted(
self,
Expand Down Expand Up @@ -970,7 +977,7 @@ def broadcast_to(
raise ak._errors.wrap_error(NotImplementedError)

def reshape(
self, x: ArrayLike, shape: tuple[int, ...], *, copy: bool | None = None
self, x: ArrayLike, shape: tuple[ShapeItem, ...], *, copy: bool | None = None
) -> TypeTracerArray:
x.touch_shape()

Expand Down Expand Up @@ -1024,9 +1031,16 @@ def nonzero(self, x: ArrayLike) -> tuple[TypeTracerArray, ...]:
try_touch_data(x)
return (TypeTracerArray._new(np.int64, (unknown_length,)),) * len(x.shape)

def where(
self, condition: ArrayLike, x1: ArrayLike, x2: ArrayLike
) -> TypeTracerArray:
condition, x1, x2 = self.broadcast_arrays(condition, x1, x2)
result_dtype = numpy.result_type(x1, x2)
return TypeTracerArray._new(result_dtype, shape=condition.shape)

def unique_values(self, x: ArrayLike) -> TypeTracerArray:
try_touch_data(x)
return TypeTracerArray._new(x.dtype, shape=(None,))
return TypeTracerArray._new(x.dtype, shape=(unknown_length,))

def concat(self, arrays, *, axis: int | None = 0) -> TypeTracerArray:
if axis is None:
Expand Down Expand Up @@ -1122,20 +1136,20 @@ def logical_and(
x2: ArrayLike,
maybe_out: ArrayLike | None = None,
) -> TypeTracerArray:
return self._apply_ufunc(numpy.sqrt, x1, x2)
return self._apply_ufunc(numpy.logical_and, x1, x2)

def logical_or(
self,
x1: ArrayLike,
x2: ArrayLike,
maybe_out: ArrayLike | None = None,
) -> TypeTracerArray:
return self._apply_ufunc(numpy.sqrt, x1, x2)
return self._apply_ufunc(numpy.logical_or, x1, x2)

def logical_not(
self, x: ArrayLike, maybe_out: ArrayLike | None = None
) -> TypeTracerArray:
return self._apply_ufunc(numpy.sqrt, x)
return self._apply_ufunc(numpy.logical_not, x)

def sqrt(self, x: ArrayLike, maybe_out: ArrayLike | None = None) -> TypeTracerArray:
return self._apply_ufunc(numpy.sqrt, x)
Expand All @@ -1149,7 +1163,7 @@ def divide(
x2: ArrayLike,
maybe_out: ArrayLike | None = None,
) -> TypeTracerArray:
return self._apply_ufunc(numpy.sqrt, x1, x2)
return self._apply_ufunc(numpy.divide, x1, x2)

############################ almost-ufuncs

Expand Down
32 changes: 22 additions & 10 deletions src/awkward/operations/ak_almost_equal.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,9 @@ def almost_equal(
The relative difference (`rtol * abs(b)`) and the absolute difference `atol`
are added together to compare against the absolute difference between `left`
and `right`.
TypeTracer arrays are not supported, as there is very little information to
be compared.
"""
left_behavior = behavior_of(left)
right_behavior = behavior_of(right)
Expand All @@ -48,6 +51,12 @@ def almost_equal(
right = to_layout(right, allow_record=False).to_packed()

backend = backend_of(left, right)
if not backend.nplike.known_data:
raise wrap_error(
NotImplementedError(
"Awkward Arrays with typetracer backends cannot yet be compared with `ak.almost_equal`."
)
)

def is_approx_dtype(left, right) -> bool:
if not dtype_exact:
Expand Down Expand Up @@ -88,25 +97,28 @@ def visitor(left, right) -> bool:
return False

if left.is_list:
return (
backend.index_nplike.array_equal(left.starts, right.starts)
and backend.index_nplike.array_equal(left.stops, right.stops)
and visitor(
left.content[: left.stops[-1]], right.content[: right.stops[-1]]
)
return backend.index_nplike.array_equal(
left.offsets, right.offsets
) and visitor(
left.content[: left.offsets[-1]], right.content[: right.offsets[-1]]
)
elif left.is_regular:
return (left.size == right.size) and visitor(left.content, right.content)
elif left.is_numpy:
return is_approx_dtype(left.dtype, right.dtype) and backend.nplike.all(
backend.nplike.isclose(
left.data, right.data, rtol=rtol, atol=atol, equal_nan=False
return (
is_approx_dtype(left.dtype, right.dtype)
and backend.nplike.all(
backend.nplike.isclose(
left.data, right.data, rtol=rtol, atol=atol, equal_nan=False
)
)
and left.shape == right.shape
)

elif left.is_option:
return backend.index_nplike.array_equal(
left.index.data < 0, right.index.data < 0
) and visitor(left.project(), right.project())
) and visitor(left.project().to_packed(), right.project().to_packed())
elif left.is_union:
return (len(left.contents) == len(right.contents)) and all(
[
Expand Down
Loading

0 comments on commit 9fcae19

Please sign in to comment.