support for cdist (#29129)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/29129
cdist(x1, x2) does the following:
- assume x1, x2 are 2-dimensional. Then x1, x2 are each considered to be
a list of vectors.
- The operation returns a matrix that is the pairwise distance between
each vector in x1 and each vector in x2. The matrix has first dimension
size equal to the number of vectors in x1 and second dimension size equal
to the number of vectors in x2.
- cdist also supports arbitrary left-hand broadcastable batch
dimensions. In this case, x1 and x2 are each considered to be a batch
of a list of vectors.
The above leads to the following name inference rule for cdist:
- In the 2D case, propagate x1.names[-2] and x2.names[-1] (because
the final result has size (x1.size[-2], x2.size[-2]).
- in the ND case, unify all the batch dimensions together to produce the
output batch dimensions and then apply the rule for the 2D case.
Furthermore, I moved all of the name checking in the implementation to
occur before name inference because name inference assumes that the
shapes are valid.
Test Plan: - new test: `pytest test/test_namedtensor.py -v -k "cdist"`
Differential Revision: D18311867
Pulled By: zou3519
fbshipit-source-id: 713d7cdda93c8fe92e7f1bd7f7c5c6e20a8138e3
