提交 febb92ea authored 作者: Brandon T. Willard's avatar Brandon T. Willard 提交者: Brandon T. Willard

Set the output dtype for Numba RandomVariable

上级 9ad68dfa
...@@ -2087,8 +2087,10 @@ def {bcast_fn_name}({bcast_fn_input_names}): ...@@ -2087,8 +2087,10 @@ def {bcast_fn_name}({bcast_fn_input_names}):
) )
# Now, create a Numba JITable function that implements the `size` parameter # Now, create a Numba JITable function that implements the `size` parameter
out_dtype = node.outputs[1].type.numpy_dtype
random_fn_global_env = { random_fn_global_env = {
bcast_fn_name: bcast_fn, bcast_fn_name: bcast_fn,
"out_dtype": out_dtype,
} }
if tuple_size > 0: if tuple_size > 0:
...@@ -2096,7 +2098,7 @@ def {bcast_fn_name}({bcast_fn_input_names}): ...@@ -2096,7 +2098,7 @@ def {bcast_fn_name}({bcast_fn_input_names}):
f""" f"""
size = to_fixed_tuple(size, tuple_size) size = to_fixed_tuple(size, tuple_size)
data = np.empty(size) data = np.empty(size, dtype=out_dtype)
for i in np.ndindex(size[:size_dims]): for i in np.ndindex(size[:size_dims]):
data[i] = {bcast_fn_name}({bcast_fn_input_names}) data[i] = {bcast_fn_name}({bcast_fn_input_names})
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论