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

Merge pull request #24932 from hawkinsp:gather

Browse Source 
PiperOrigin-RevId: 697632284
This commit is contained in:
jax authors 2024-11-18 08:14:23 -08:00
commit 65f9c7855e
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/lax/lax.py
View File

@ -4544,6 +4544,8 @@ def shape_as_value(shape: core.Shape):
"""Converts a shape that may contain Poly values into a JAX value."""
if len(shape) == 0:
return full((0,), np.array(0, np.int64))
if core.is_constant_shape(shape):
return np.asarray(shape, dtype=np.int64)
dims = [
expand_dims(convert_element_type(core.dimension_as_value(d), np.int64),
(0,))