Skip to content

[Relax][PyTorch] Add boolean tensor support for max operation and corresponding test case #18530

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mshr-h merged 4 commits into from
Nov 29, 2025
Merged

[Relax][PyTorch] Add boolean tensor support for max operation and corresponding test case #18530

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
3600910
[Relax][PyTorch] Implement boolean tensor support for max operation a…
mshr-h Nov 29, 2025
049d85b
sdpa shouldn't be decomposed as tvm has the corresponding relax op
mshr-h Nov 29, 2025
343432e
fix isin ops test
mshr-h Nov 29, 2025
53c0949
Merge branch 'main' into fix-any-bool
mshr-h Nov 29, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions python/tvm/relax/frontend/torch/base_fx_graph_translator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1634,6 +1634,12 @@ def _any(self, node: fx.Node) -> relax.Var:
dim = args[1] if len(node.args) > 1 else node.kwargs.get("dim", None)
keepdim = args[2] if len(node.args) > 2 else node.kwargs.get("keepdim", False)

# max doesn't support boolean tensors directly, so we compute it in int8 and cast back
if x.struct_info.dtype == "bool":
x = relax.op.astype(x, "int8")
ret = relax.op.max(x, dim, keepdims=keepdim)
return self.block_builder.emit(relax.op.astype(ret, "bool"))

# For boolean tensors, any is equivalent to max (checking if any element is True)
return self.block_builder.emit(relax.op.max(x, dim, keepdims=keepdim))

Expand Down
176 changes: 51 additions & 125 deletions tests/python/relax/test_frontend_from_exported_program.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1693,8 +1693,10 @@ def main(
with R.dataflow():
lv: R.Tensor((10, 10, 1), dtype="float32") = R.reshape(x, R.shape([10, 10, 1]))
lv1: R.Tensor((10, 10, 8), dtype="bool") = R.equal(lv, test_elements)
lv2: R.Tensor((10, 10), dtype="bool") = R.max(lv1, axis=[-1], keepdims=False)
gv: R.Tuple(R.Tensor((10, 10), dtype="bool")) = (lv2,)
lv2: R.Tensor((10, 10, 8), dtype="int8") = R.astype(lv1, dtype="int8")
lv3: R.Tensor((10, 10), dtype="int8") = R.max(lv2, axis=[-1], keepdims=False)
lv4: R.Tensor((10, 10), dtype="bool") = R.astype(lv3, dtype="bool")
gv: R.Tuple(R.Tensor((10, 10), dtype="bool")) = (lv4,)
R.output(gv)
return gv

Expand Down Expand Up @@ -4118,71 +4120,22 @@ def main(
v: R.Tensor((32, 8, 128, 64), dtype="float32"),
) -> R.Tuple(R.Tensor((32, 8, 128, 64), dtype="float32")):
with R.dataflow():
lv: R.Tensor((32, 8, 128, 64), dtype="float32") = R.multiply(
q, R.const(0.35355338454246521, "float32")
lv: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
q, axes=[0, 2, 1, 3]
)
lv1: R.Tensor((32, 8, 64, 128), dtype="float32") = R.permute_dims(
k, axes=[0, 1, 3, 2]
lv1: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
k, axes=[0, 2, 1, 3]
)
lv2: R.Tensor((32, 8, 64, 128), dtype="float32") = R.multiply(
lv1, R.const(0.35355338454246521, "float32")
lv2: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
v, axes=[0, 2, 1, 3]
)
lv3: R.Tensor((32, 8, 128, 64), dtype="float32") = R.broadcast_to(
lv, R.shape([32, 8, 128, 64])
lv3: R.Tensor((32, 128, 8, 64), dtype="float32") = R.nn.attention(
lv, lv1, lv2, scale=None, causal_mask=None, window_size=None
)
lv4: R.Tensor((256, 128, 64), dtype="float32") = R.reshape(
lv3, R.shape([256, 128, 64])
lv4: R.Tensor((32, 8, 128, 64), dtype="float32") = R.permute_dims(
lv3, axes=[0, 2, 1, 3]
)
lv5: R.Tensor((32, 8, 64, 128), dtype="float32") = R.broadcast_to(
lv2, R.shape([32, 8, 64, 128])
)
lv6: R.Tensor((256, 64, 128), dtype="float32") = R.reshape(
lv5, R.shape([256, 64, 128])
)
lv7: R.Tensor((256, 128, 128), dtype="float32") = R.matmul(
lv4, lv6, out_dtype="float32"
)
lv8: R.Tensor((32, 8, 128, 128), dtype="float32") = R.reshape(
lv7, R.shape([32, 8, 128, 128])
)
lv9: R.Tensor((32, 8, 128, 128), dtype="float32") = R.nn.softmax(lv8, axis=-1)
lv10: R.Tensor((32, 8, 128, 128), dtype="bool") = R.equal(
lv8, R.const(float("-inf"), "float32")
)
lv11: R.Tensor((32, 8, 128, 128), dtype="bool") = R.logical_not(lv10)
lv12: R.Tensor((32, 8, 128, 1), dtype="bool") = R.max(
lv11, axis=[-1], keepdims=True
)
lv13: R.Tensor((32, 8, 128, 1), dtype="bool") = R.logical_not(lv12)
lv14: R.Tensor((32, 8, 128, 128), dtype="float32") = R.full_like(
lv9, R.const(0, "int32"), dtype="void"
)
lv15: R.Tensor((32, 8, 128, 128), dtype="float32") = R.where(lv13, lv14, lv9)
lv16: R.Tensor((32, 8, 128, 128), dtype="float32") = R.broadcast_to(
lv15, R.shape([32, 8, 128, 128])
)
lv17: R.Tensor((256, 128, 128), dtype="float32") = R.reshape(
lv16, R.shape([256, 128, 128])
)
lv18: R.Tensor((32, 8, 128, 64), dtype="float32") = R.broadcast_to(
v, R.shape([32, 8, 128, 64])
)
lv19: R.Tensor((256, 128, 64), dtype="float32") = R.reshape(
lv18, R.shape([256, 128, 64])
)
lv20: R.Tensor((256, 128, 64), dtype="float32") = R.matmul(
lv17, lv19, out_dtype="float32"
)
lv21: R.Tensor((32, 8, 128, 64), dtype="float32") = R.reshape(
lv20, R.shape([32, 8, 128, 64])
)
lv22: R.Tensor((128, 32, 8, 64), dtype="float32") = R.permute_dims(
lv21, axes=[2, 0, 1, 3]
)
lv23: R.Tensor((32, 8, 128, 64), dtype="float32") = R.permute_dims(
lv22, axes=[1, 2, 0, 3]
)
gv: R.Tuple(R.Tensor((32, 8, 128, 64), dtype="float32")) = (lv23,)
gv: R.Tuple(R.Tensor((32, 8, 128, 64), dtype="float32")) = (lv4,)
R.output(gv)
return gv

Expand All @@ -4200,72 +4153,22 @@ def main(
mask: R.Tensor((32, 8, 128, 128), dtype="float32"),
) -> R.Tuple(R.Tensor((32, 8, 128, 64), dtype="float32")):
with R.dataflow():
lv: R.Tensor((32, 8, 128, 64), dtype="float32") = R.multiply(
q, R.const(0.35355338454246521, "float32")
)
lv1: R.Tensor((32, 8, 64, 128), dtype="float32") = R.permute_dims(
k, axes=[0, 1, 3, 2]
)
lv2: R.Tensor((32, 8, 64, 128), dtype="float32") = R.multiply(
lv1, R.const(0.35355338454246521, "float32")
)
lv3: R.Tensor((32, 8, 128, 64), dtype="float32") = R.broadcast_to(
lv, R.shape([32, 8, 128, 64])
)
lv4: R.Tensor((256, 128, 64), dtype="float32") = R.reshape(
lv3, R.shape([256, 128, 64])
)
lv5: R.Tensor((32, 8, 64, 128), dtype="float32") = R.broadcast_to(
lv2, R.shape([32, 8, 64, 128])
)
lv6: R.Tensor((256, 64, 128), dtype="float32") = R.reshape(
lv5, R.shape([256, 64, 128])
)
lv7: R.Tensor((256, 128, 128), dtype="float32") = R.matmul(
lv4, lv6, out_dtype="float32"
)
lv8: R.Tensor((32, 8, 128, 128), dtype="float32") = R.reshape(
lv7, R.shape([32, 8, 128, 128])
lv: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
q, axes=[0, 2, 1, 3]
)
lv9: R.Tensor((32, 8, 128, 128), dtype="float32") = R.add(lv8, mask)
lv10: R.Tensor((32, 8, 128, 128), dtype="float32") = R.nn.softmax(lv9, axis=-1)
lv11: R.Tensor((32, 8, 128, 128), dtype="bool") = R.equal(
lv9, R.const(float("-inf"), "float32")
lv1: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
k, axes=[0, 2, 1, 3]
)
lv12: R.Tensor((32, 8, 128, 128), dtype="bool") = R.logical_not(lv11)
lv13: R.Tensor((32, 8, 128, 1), dtype="bool") = R.max(
lv12, axis=[-1], keepdims=True
lv2: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
v, axes=[0, 2, 1, 3]
)
lv14: R.Tensor((32, 8, 128, 1), dtype="bool") = R.logical_not(lv13)
lv15: R.Tensor((32, 8, 128, 128), dtype="float32") = R.full_like(
lv10, R.const(0, "int32"), dtype="void"
lv3: R.Tensor((32, 128, 8, 64), dtype="float32") = R.nn.attention_bias(
lv, lv1, lv2, mask, scale=None, causal_mask=None, window_size=None
)
lv16: R.Tensor((32, 8, 128, 128), dtype="float32") = R.where(lv14, lv15, lv10)
lv17: R.Tensor((32, 8, 128, 128), dtype="float32") = R.broadcast_to(
lv16, R.shape([32, 8, 128, 128])
lv4: R.Tensor((32, 8, 128, 64), dtype="float32") = R.permute_dims(
lv3, axes=[0, 2, 1, 3]
)
lv18: R.Tensor((256, 128, 128), dtype="float32") = R.reshape(
lv17, R.shape([256, 128, 128])
)
lv19: R.Tensor((32, 8, 128, 64), dtype="float32") = R.broadcast_to(
v, R.shape([32, 8, 128, 64])
)
lv20: R.Tensor((256, 128, 64), dtype="float32") = R.reshape(
lv19, R.shape([256, 128, 64])
)
lv21: R.Tensor((256, 128, 64), dtype="float32") = R.matmul(
lv18, lv20, out_dtype="float32"
)
lv22: R.Tensor((32, 8, 128, 64), dtype="float32") = R.reshape(
lv21, R.shape([32, 8, 128, 64])
)
lv23: R.Tensor((128, 32, 8, 64), dtype="float32") = R.permute_dims(
lv22, axes=[2, 0, 1, 3]
)
lv24: R.Tensor((32, 8, 128, 64), dtype="float32") = R.permute_dims(
lv23, axes=[1, 2, 0, 3]
)
gv: R.Tuple(R.Tensor((32, 8, 128, 64), dtype="float32")) = (lv24,)
gv: R.Tuple(R.Tensor((32, 8, 128, 64), dtype="float32")) = (lv4,)
R.output(gv)
return gv

Expand All @@ -4278,7 +4181,7 @@ def main(
),
{},
Expected1,
run_ep_decomposition=True,
run_ep_decomposition=False,
Copy link
Copy Markdown
Member

@tlopex tlopex Nov 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just want to know why should run_ep_decomposition be disabled here?
Is the intention that Relax already supports this composite op, so decomposition is unnecessary?

Copy link
Copy Markdown
Contributor Author

@mshr-h mshr-h Nov 29, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is the intention that Relax already supports this composite op, so decomposition is unnecessary?

Exactly. I'm thinking we should provide a way for users to control the decomposition. Or we have to define list of ops that we want not to be decomposed.

)

verify_model(
Expand All @@ -4291,7 +4194,7 @@ def main(
),
{},
Expected2,
run_ep_decomposition=True,
run_ep_decomposition=False,
)

# Test 2D input (seq_len, head_dim) - bug fix for #18441
Expand Down Expand Up @@ -7307,6 +7210,29 @@ def main(
verify_model(Take(), example_args, {}, Expected)


def test_any():
class AnyAten(torch.nn.Module):
def forward(self, x):
return torch.ops.aten.any(x, dim=1)

@tvm.script.ir_module
class Expected:
@R.function
def main(
x: R.Tensor((2, 3), dtype="bool"),
) -> R.Tuple(R.Tensor((2,), dtype="bool")):
with R.dataflow():
lv: R.Tensor((2, 3), dtype="int8") = relax.op.astype(x, dtype="int8")
lv2: R.Tensor((2,), dtype="int8") = relax.op.max(lv, axis=1, keepdims=False)
lv3: R.Tensor((2,), dtype="bool") = relax.op.astype(lv2, dtype="bool")
gv: R.Tuple(R.Tensor((2,), dtype="bool")) = (lv3,)
R.output(gv)
return gv

example_args = (torch.tensor([[0, 0, 0], [0, 1, 0]], dtype=torch.bool),)
verify_model(AnyAten(), example_args, {}, Expected)


def test_std():
class Std(Module):
def forward(self, x):
Expand Down
Loading