Merge pull request #65 from Quansight-Labs/where_et_al

wrap more functions to complete a minimal viable feature set

Author: ev-br

autogen/numpy_api_dump.py

@@ -172,10 +172,6 @@ def choose(a, choices, out=None, mode="raise"):
     raise NotImplementedError
 
 
-def clip(a, a_min, a_max, out=None, **kwargs):
-    raise NotImplementedError
-
-
 def common_type(*arrays):
     raise NotImplementedError
 
@@ -216,22 +212,6 @@ def deprecate_with_doc(msg):
     raise NotImplementedError
 
 
-def diag_indices(n, ndim=2):
-    raise NotImplementedError
-
-
-def diag_indices_from(arr):
-    raise NotImplementedError
-
-
-def diagflat(v, k=0):
-    raise NotImplementedError
-
-
-def diagonal(a, offset=0, axis1=0, axis2=1):
-    raise NotImplementedError
-
-
 def digitize(x, bins, right=False):
     raise NotImplementedError
 
@@ -240,10 +220,6 @@ def disp(mesg, device=None, linefeed=True):
     raise NotImplementedError
 
 
-def dot(a, b, out=None):
-    raise NotImplementedError
-
-
 def ediff1d(ary, to_end=None, to_begin=None):
     raise NotImplementedError
 
@@ -260,10 +236,6 @@ def extract(condition, arr):
     raise NotImplementedError
 
 
-def fill_diagonal(a, val, wrap=False):
-    raise NotImplementedError
-
-
 def find_common_type(array_types, scalar_types):
     raise NotImplementedError
 
@@ -393,14 +365,6 @@ def in1d(ar1, ar2, assume_unique=False, invert=False):
     raise NotImplementedError
 
 
-def indices(dimensions, dtype=int, sparse=False):
-    raise NotImplementedError
-
-
-def inner(a, b, /):
-    raise NotImplementedError
-
-
 def insert(arr, obj, values, axis=None):
     raise NotImplementedError
 
@@ -417,10 +381,6 @@ def is_busday(dates, weekmask="1111100", holidays=None, busdaycal=None, out=None
     raise NotImplementedError
 
 
-def isclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False):
-    raise NotImplementedError
-
-
 def isfortran(a):
     raise NotImplementedError
 
@@ -437,14 +397,6 @@ def issctype(rep):
     raise NotImplementedError
 
 
-def issubclass_(arg1, arg2):
-    raise NotImplementedError
-
-
-def issubdtype(arg1, arg2):
-    raise NotImplementedError
-
-
 def issubsctype(arg1, arg2):
     raise NotImplementedError
 
@@ -624,10 +576,6 @@ def obj2sctype(rep, default=None):
     raise NotImplementedError
 
 
-def outer(a, b, out=None):
-    raise NotImplementedError
-
-
 def packbits(a, /, axis=None, bitorder="big"):
     raise NotImplementedError
 
@@ -696,10 +644,6 @@ def putmask(a, mask, values):
     raise NotImplementedError
 
 
-def ravel(a, order="C"):
-    raise NotImplementedError
-
-
 def recfromcsv(fname, **kwargs):
     raise NotImplementedError
 
@@ -754,10 +698,6 @@ def sctype2char(sctype):
     raise NotImplementedError
 
 
-def searchsorted(a, v, side="left", sorter=None):
-    raise NotImplementedError
-
-
 def select(condlist, choicelist, default=0):
     raise NotImplementedError
 
@@ -811,22 +751,10 @@ def show():
     raise NotImplementedError
 
 
-def sinc(x):
-    raise NotImplementedError
-
-
-def sometrue(*args, **kwargs):
-    raise NotImplementedError
-
-
 def sort_complex(a):
     raise NotImplementedError
 
 
-def swapaxes(a, axis1, axis2):
-    raise NotImplementedError
-
-
 def take(a, indices, axis=None, out=None, mode="raise"):
     raise NotImplementedError
 
@@ -835,10 +763,6 @@ def tensordot(a, b, axes=2):
     raise NotImplementedError
 
 
-def trace(a, offset=0, axis1=0, axis2=1, dtype=None, out=None):
-    raise NotImplementedError
-
-
 def trapz(y, x=None, dx=1.0, axis=-1):
     raise NotImplementedError
 
@@ -855,18 +779,6 @@ def union1d(ar1, ar2):
     raise NotImplementedError
 
 
-def unique(
-    ar,
-    return_index=False,
-    return_inverse=False,
-    return_counts=False,
-    axis=None,
-    *,
-    equal_nan=True,
-):
-    raise NotImplementedError
-
-
 def unpackbits(a, /, axis=None, count=None, bitorder="big"):
     raise NotImplementedError
 
@@ -875,13 +787,5 @@ def unwrap(p, discont=None, axis=-1, *, period=6.283185307179586):
     raise NotImplementedError
 
 
-def vdot(a, b, /):
-    raise NotImplementedError
-
-
-def where(condition, x, y, /):
-    raise NotImplementedError
-
-
 def who(vardict=None):
     raise NotImplementedError

torch_np/_detail/_flips.py

@@ -29,6 +29,8 @@ def rot90(m_tensor, k=1, axes=(0, 1)):
 
 
 def swapaxes(tensor, axis1, axis2):
+    axis1 = _util.normalize_axis_index(axis1, tensor.ndim)
+    axis2 = _util.normalize_axis_index(axis2, tensor.ndim)
     return torch.swapaxes(tensor, axis1, axis2)
 
 

torch_np/_detail/_reductions.py

@@ -143,11 +143,19 @@ def mean(tensor, axis=None, dtype=None, *, where=NoValue):
 
     dtype = _atleast_float(dtype, tensor.dtype)
 
+    is_half = dtype == torch.float16
+    if is_half:
+        # XXX revisit when the pytorch version has pytorch/pytorch#95166
+        dtype = torch.float32
+
     if axis is None:
         result = tensor.mean(dtype=dtype)
     else:
         result = tensor.mean(dtype=dtype, dim=axis)
 
+    if is_half:
+        result = result.to(torch.float16)
+
     return result
 
 

torch_np/_detail/_util.py

@@ -36,7 +36,7 @@ class UFuncTypeError(TypeError, RuntimeError):
 
 def cast_if_needed(tensor, dtype):
     # NB: no casting if dtype=None
-    if tensor.dtype != dtype:
+    if dtype is not None and tensor.dtype != dtype:
         tensor = tensor.to(dtype)
     return tensor