I see your point and it seems fair.
In the current implementation, one solution I can think of is leveraging the
`check` function:
```python
import tvm
from tvm import relay
from tvm.relay.dataflow_pattern import *
def check(pre):
return (pre.args[0].checked_type.dtype == 'float32' and
pre.args[1].checked_type.dtype == 'float32')
pat = is_op('add')(wildcard(), wildcard())
x = relay.var('x', shape=(10, 10), dtype='float32')
out = relay.add(x, x)
func = relay.Function([x], out)
mod = tvm.IRModule()
mod['main'] = func
mod = relay.transform.InferType()(mod)
print(pat.partition(mod['main'].body, check=check))
```
In short, you can implement a check function which does any forms of checking
by looking into the matched subgraph.
In case you only need to know if matching or not and do not want to partition
the graph, you can use `rewrite` to mimic the above functionality:
```python
import tvm
from tvm import relay
from tvm.relay.dataflow_pattern import *
class MyCallback(DFPatternCallback):
def __init__(self):
self.in1 = wildcard()
self.in2 = wildcard()
self.pattern = is_op('add')(self.in1, self.in2)
self.match = False
def callback(self, pre, post, node_map):
if (node_map[self.in1][0].checked_type.dtype == 'float32' and
node_map[self.in2][0].checked_type.dtype == 'float32'):
self.match = True
return post
x = relay.var('x', shape=(10, 10), dtype='float32')
out = relay.add(x, x)
func = relay.Function([x], out)
mod = tvm.IRModule()
mod['main'] = func
mod = relay.transform.InferType()(mod)
callback = MyCallback()
rewrite(callback, mod['main'].body)
print(callback.match)
```
When matching the pattern, `rewrite` will call the callback function for graph
mutation. You can of course add any checks here and maintain your own "match"
status.
While the above solutions are working, there are definitely imperfect,
especially when the pattern is complex. In long term, we may want to support
partial type matching in the pattern language.
cc @mbrookhart
---
[Visit
Topic](https://discuss.tvm.ai/t/pattenlang-how-to-match-op-according-to-element-type-of-input-and-output/6846/6)
to respond.
You are receiving this because you enabled mailing list mode.
To unsubscribe from these emails, [click
here](https://discuss.tvm.ai/email/unsubscribe/718dc81d65775478e665cb3058af7b1e63259f83066a372388f9677b39fec672).