Is there Relay instruction 'pytorch nn.embedding'?

choi95 · May 22, 2022, 10:03am

by any chance, relay.nn.sparse_dense is same to nn.embedding?

comaniac · May 23, 2022, 12:03am

nn.embedding is corresponding to take op in relay. See:

apache/tvm/blob/main/python/tvm/relay/frontend/pytorch.py#L2058


    else:
        data0, data1 = self.pytorch_promote_types(inputs, input_types)
        return get_relay_op("subtract")(data0, data1)


def rsub(self, inputs, input_types):
    data0, data1, alpha = self.pytorch_promote_types(inputs, input_types)


    # note: rsub means data0 and data1 swap places
    return get_relay_op("subtract")(data1, alpha * data0)


def embedding(self, inputs, input_types):
    weight = inputs[0]
    indices = inputs[1]


    return _op.take(weight, indices.astype("int32"), axis=0)


def one_hot(self, inputs, input_types):
    indices = inputs[0].astype("int32")
    num_classes = inputs[1]
    if num_classes == -1:
        msg = "Inferring the number of classes is not yet supported."