mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #10385 from jakevdp:promote-types-comment

Browse Source 
PiperOrigin-RevId: 443116925
This commit is contained in:
jax authors 2022-04-20 09:39:46 -07:00
commit 606d1c2402
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
jax/_src/dtypes.py
View File

@ -327,6 +327,8 @@ def promote_types(a, b):
Returns:
A :class:`numpy.dtype` object.
"""
# Note: we deliberately avoid `if a in _weak_types` here because we want to check
# object identity, not object equality, due to the behavior of np.dtype.__eq__
a = a if any(a is t for t in _weak_types) else np.dtype(a)
b = b if any(b is t for t in _weak_types) else np.dtype(b)
return np.dtype(_least_upper_bound(a, b))