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

make lax.full require concrete shapes

Browse Source 
improves error message for #204
This commit is contained in:
Matthew Johnson 2019-01-07 12:28:52 -08:00
parent 5031016465
commit df87d5ce43
1 changed files with 10 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
jax/lax.py
View File

@ -418,6 +418,13 @@ def tie_in(x, y):
return tie_in_p.bind(x, y)
def full(shape, fill_value, dtype):
try:
shape = tuple(map(int, shape))
except TypeError:
msg = ("`full` requires shapes to be concrete. If using `jit`, try using "
"`static_argnums` or applying `jit` to smaller subfunctions instead.")
raise TypeError(msg)
if onp.shape(fill_value):
msg = "full must be called with scalar fill_value, got fill_value.shape {}."
raise TypeError(msg.format(onp.shape(fill_value)))
@ -2532,7 +2539,9 @@ def _check_shapelike(fun_name, arg_name, obj):
def _dynamic_slice_indices(operand, start_indices):
if isinstance(start_indices, (tuple, list)):
start_indices = concatenate([reshape(i, [1]) for i in start_indices], 0)
return rem(start_indices, onp.array(operand.shape, start_indices.dtype))
# map int over operand.shape to raise any dynamic-shape errors
shape = onp.asarray(map(int, operand.shape), start_indices.dtype)
return rem(start_indices, shape)
def _const(example, val):