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

Removed unused _upcast_fp16_for_computation

Browse Source 
PiperOrigin-RevId: 691409888
This commit is contained in:
Sergei Lebedev 2024-10-30 07:23:08 -07:00 committed by jax authors
parent a45b0856c5
commit 2b70ad30fb
1 changed files with 0 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
jax/_src/lax/lax.py
View File

@ -1912,17 +1912,6 @@ def reciprocal(x: ArrayLike) -> Array:
r"""Elementwise reciprocal: :math:`1 \over x`."""
return integer_pow(x, -1)
def _upcast_fp16_for_computation(f):
@functools.wraps(f)
def f_wrapped(x):
dtype = _dtype(x)
if dtype == np.float16 or dtype == dtypes.bfloat16:
return convert_element_type(
f(convert_element_type(x, np.float32)), dtype)
return f(x)
return f_wrapped
def tan(x: ArrayLike) -> Array:
r"""Elementwise tangent: :math:`\mathrm{tan}(x)`."""
return tan_p.bind(x)