提交 53763f5d authored 作者: Ricardo Vieira's avatar Ricardo Vieira 提交者: Ricardo Vieira

Don't use objectmode with vector boolean `inc_subtensor`

上级 67017a66
......@@ -130,7 +130,6 @@ def numba_funcify_AdvancedSubtensor(op, node, **kwargs):
if isinstance(idx.type, TensorType)
]
# Special case for consecutive consecutive vector indices
def broadcasted_to(x_bcast: tuple[bool, ...], to_bcast: tuple[bool, ...]):
# Check that x is not broadcasted to y based on broadcastable info
if len(x_bcast) < len(to_bcast):
......@@ -176,7 +175,14 @@ def numba_funcify_AdvancedSubtensor(op, node, **kwargs):
or (
isinstance(op, AdvancedIncSubtensor)
and not op.set_instead_of_inc
and not op.ignore_duplicates
and not (
op.ignore_duplicates
# Only vector integer indices can have "duplicates", not scalars or boolean vectors
or all(
adv_idx["ndim"] == 0 or adv_idx["dtype"] == "bool"
for adv_idx in adv_idxs
)
)
)
):
return generate_fallback_impl(op, node, **kwargs)
......
......@@ -314,8 +314,16 @@ def test_AdvancedIncSubtensor1(x, y, indices):
np.arange(3 * 4 * 5).reshape((3, 4, 5)),
-np.arange(1 * 4 * 5).reshape(1, 4, 5),
(np.array([True, False, False])), # Broadcasted boolean index
False, # It shouldn't matter what we set this to, boolean indices cannot be duplicate
False,
False,
),
(
np.arange(3 * 4 * 5).reshape((3, 4, 5)),
-np.arange(1 * 4 * 5).reshape(1, 4, 5),
(np.array([True, False, False])), # Broadcasted boolean index
True, # It shouldn't matter what we set this to, boolean indices cannot be duplicate
False,
False,
),
(
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论