[Relay][Op] Add compute, schedule, and tests for expand_dims and squeeze #2133
Conversation
|
Incidentally, is there a fundamental reason that Relay's squeeze needs to have a different calling convention from NNVM's? Namely, taking a list instead of a tuple and not allowing a single int in place of a list of one |
|
Also is there a reason Relay's |
|
topi squeeze doesn't handle nonpositive dimensions. implementing those properly will require some thought, because resolving them requires shape computation |
|
the single int in place of a list could be handled by the register compute function |
|
I see, regarding the shape computation issue (I am not sure I agree with the rationale since many other Relay functions do the same thing with axis args; we should be consistent). I don't think the register compute function alone can or should handle the discrepancy in calling conventions because the Relay function |
|
It seems that we will need to add negative axis support to topi, to make it compatible. Note that we only need to know the dimension of the shape and they are known in topi compute. We can provide the support for passing a single int in the python wrapping, by detecting the input type. |
| check_binary_op(opfunc, ref) | ||
|
|
||
|
|
||
| def test_expand_dims(): |
There was a problem hiding this comment.
Invoke this function from __main__
|
@tqchen thanks for the explanation. If I understand correctly, couldn't we handle negative entries in |
|
This is a simple feature that seems would make topi more general, so we could do it in the topi side. Also given that this is already a simple function, our goal is to move most of the compute to c++ |
|
We should update the relation then, right? |
|
@joshpoll yes |
slyubomirsky
force-pushed
the
squeeze-expand-dims
branch
from
ebe571b to
aace0d3
Compare
| # expand_dims | ||
| @register_compute("expand_dims") | ||
| def expand_dims_compiler(attrs, inputs, output_type, target): | ||
| """Compiler for expand_dims.""" |
There was a problem hiding this comment.
make sure expand_dims in topi support negative index(I think they may already)
| # squeeze | ||
| @register_compute("squeeze") | ||
| def squeeze_compiler(attrs, inputs, output_type, target): | ||
| """Compiler for squeeze dims.""" |
There was a problem hiding this comment.
squeeze already support negative index, and None I think so you just have to redirect to topi call
There was a problem hiding this comment.
I think we should rename this to "compute" to match the style in the rest of the system.
|
Thanks, @joshpoll @slyubomirsky @jroesch @siju-samuel , this is merged. |
|
see followup changes in #2163 |
6 participants
As the title states, these are compute and schedule primitives for operators
expand_dimsandsqueeze, with tests included. Please review (e.g., @jroesch, @MarisaKirisame) for correctness and style.