Question about nn.relay.batch_matmul

1

https://tvm.apache.org/docs/api/python/relay/nn.html#tvm.relay.nn.batch_matmul

Hi! I’m currently dependent on v0.0.4, and am trying to use Relay as a front-end.

So the following

    x = relay.var("x", relay.TensorType((3, 3, 3), "int32"))
    y = relay.var("y", relay.TensorType((3, 3, 3), "int32"))
    return relay.Function([x, y], relay.nn.batch_matmul(x, y))

lowers to:

fn (%x: Tensor[(3, 3, 3), int32], %y: Tensor[(3, 3, 3), int32]) -> Tensor[(3, 3, 3), int32] {
  let %x1: Tensor[(3, 3, 3), int32] = nn.batch_matmul(%x, %y) /* ty=Tensor[(3, 3, 3), int32] */;
  %x1
}

Ok great! However, I want to try a batch with different dimensions. For example, this

    x = relay.var("x", relay.TensorType((3, 2, 5), "int32"))
    y = relay.var("y", relay.TensorType((3, 5, 2), "int32"))
    return relay.Function([x, y], relay.nn.batch_matmul(x, y))

results in:

TVMError: 
Error(s) have occurred. The program has been annotated with them:
In `main`: 
v0.0.4
fn (%x: Tensor[(3, 2, 5), int32], %y: Tensor[(3, 5, 2), int32]) {
  nn.batch_matmul(%x, %y) an internal invariant was violated while typechecking your program [11:50:16] ../src/relay/op/nn/nn.cc:857: Check failed: reporter->AssertEQ(x->shape[2], y->shape[2]): BatchDot: shapes of x and y is inconsistent,  x shape=[3, 2, 5], y shape=[3, 5, 2]
; 
}

Can someone explain why this is illegal? My understanding of nn.batch_matmul may be incorrect.

2

Hello @cgyurgyik, Relay’s batch_matmul implements A[i, :, :] * B[i, : , :]^T (note the transpose). This means that both the input matrices must have the same size. The error message shapes of x and y is inconsistent, x shape=[3, 2, 5], y shape=[3, 5, 2] is telling you that. If you switch y to y = relay.var("y", relay.TensorType((3, 2, 5), "int32")), everything should work.

1 Like
3

Ah thanks! I saw the transpose as well, it obviously didn’t click :slight_smile: