Skip to content

Commit

Permalink
Raise NotImplementedError for boolean scalar indexing
Browse files Browse the repository at this point in the history
This is a special-edge case that behaves differently than other boolean indexing. It adds a new dimension (either empty or with one entry). The logic in `make_node` and `infer_shape` don't handle this.
  • Loading branch information
ricardoV94 committed Oct 31, 2024
1 parent 970543f commit b08fb2a
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 0 deletions.
5 changes: 5 additions & 0 deletions pytensor/tensor/subtensor.py
Original file line number Diff line number Diff line change
Expand Up @@ -2627,6 +2627,11 @@ def as_index_variable(idx):
idx = as_tensor_variable(idx)
if idx.type.dtype not in discrete_dtypes:
raise TypeError("index must be integers or a boolean mask")
if idx.type.dtype == "bool" and idx.type.ndim == 0:
raise NotImplementedError(
"Boolean scalar indexing not implemented. "
"Open an issue in https://github.com/pymc-devs/pytensor/issues if you need this behavior."
)
return idx


Expand Down
5 changes: 5 additions & 0 deletions tests/tensor/test_subtensor.py
Original file line number Diff line number Diff line change
Expand Up @@ -2228,6 +2228,11 @@ def fun(x, y):
mode=self.mode,
)

def test_boolean_scalar_raises(self):
x = vector("x")
with pytest.raises(NotImplementedError):
x[np.array(True)]


class TestInferShape(utt.InferShapeTester):
@staticmethod
Expand Down

0 comments on commit b08fb2a

Please sign in to comment.