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

Fix JAX functions to work if the default gather mode is set to "fill".

Browse Source 
These functions really do want "clip".
This commit is contained in:
Peter Hawkins 2021-09-30 14:21:05 -04:00
parent 4f6a6c5dec
commit f8ba024621
2 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
jax/_src/numpy/lax_numpy.py
View File

@ -5431,7 +5431,8 @@ def _take(a, indices, axis: Optional[int] = None, out=None, mode=None):
collapsed_slice_dims=(axis_idx,),
start_index_map=(axis_idx,))
return lax.gather(a, indices[..., None], dimension_numbers=dnums,
slice_sizes=tuple(slice_sizes))
slice_sizes=tuple(slice_sizes),
mode="clip")
def _normalize_index(index, axis_size):

2
jax/_src/scipy/special.py
View File

@ -1028,7 +1028,7 @@ def _sph_harm(m: jnp.ndarray,
cos_colatitude = jnp.cos(phi)
legendre = _gen_associated_legendre(n_max, cos_colatitude, True)
legendre_val = legendre[abs(m), n, jnp.arange(len(n))]
legendre_val = legendre.at[abs(m), n, jnp.arange(len(n))].get(mode="clip")
angle = abs(m) * theta
vandermonde = lax.complex(jnp.cos(angle), jnp.sin(angle))