add more clear comments, and reshape stuff
@@ -56,6 +56,7 @@ import io.improbable.keanu.vertices.dbl.nonprobabilistic.operators.unary.TakeVer
import io.improbable.keanu.vertices.dbl.nonprobabilistic.operators.unary.TanVertex;
import io.improbable.keanu.vertices.intgr.IntegerVertex;
import io.improbable.keanu.vertices.intgr.nonprobabilistic.CastToIntegerVertex;
+import org.apache.commons.lang3.ArrayUtils;
import java.util.Map;
import java.util.function.Function;
@@ -99,14 +100,16 @@ public abstract class DoubleVertex extends Vertex<DoubleTensor> implements Doubl
/**
* Matrix product of two vertices
*
- * @param that a double vertex to matrix multiply
- * @return a vertex that represents the matrix multiplication of two vertices. Note that this operation
- * does that always return a MatrixMultiplicationVertex.
- * - If both left and right operands are rank 2, they are multiplied like conventional matrices.
+ * @param that a double vertex representing a matrix or a vector to matrix multiply
+ * @return a vertex that represents the matrix multiplication of two vertices.
* - If both left and right operands are rank 1, they are promoted to a matrix by prepending a 1 to its dimensions.
* After matrix multiplication, it is reshaped to be a scalar. This is essentially a dot product.
+ * This returns a ReshapeVertex.
* - If only one of the operands are rank 1, it is promoted to a matrix by prepending a 1 to its dimensions.
* After matrix multiplication, the appended 1 is removed. This is essentially a matrix-vector product.
+ * This returns a ReshapeVertex.
+ * - Otherwise, they are multiplied like conventional matrices.
+ * This returns a MatrixMultiplicationVertex.
*/
public DoubleVertex matrixMultiply(DoubleVertex that) {
int leftRank = this.getRank();
@@ -120,9 +123,9 @@ public abstract class DoubleVertex extends Vertex<DoubleTensor> implements Doubl
if (leftRank == 1 && rightRank == 1) {
return result.reshape();
} else if (leftRank == 1) {
- return result.reshape(result.getShape()[1]);
+ return result.reshape(ArrayUtils.remove(result.getShape(), 0));
} else if (rightRank == 1) {
- return result.reshape(result.getShape()[0]);
+ return result.reshape(ArrayUtils.remove(result.getShape(), 1));
} else {
return result;
}
