Add matrix multiplication operator for torch_tensor (#1379)

* Add matrix multiplication operator for torch_tensor

* Add conditional namespace

* add test case

---------

Co-authored-by: Daniel Falbel <[email protected]>

Author: caiohamamura

NAMESPACE

@@ -877,6 +877,7 @@ export(with_enable_grad)
 export(with_no_grad)
 export(with_torch_manual_seed)
 export(yield)
+if (getRversion() >= "4.3.0") S3method("%*%", torch_tensor)
 importFrom(Rcpp,sourceCpp)
 importFrom(bit64,as.integer64)
 importFrom(callr,r)

R/operators.R

@@ -175,6 +175,14 @@
   torch_logical_not(x)
 }
 
+#' @rawNamespace if (getRversion() >= "4.3.0") S3method("%*%", torch_tensor)
+`%*%.torch_tensor` <- function(e1, e2) {
+  if (!is_torch_tensor(e2)) {
+    e2 <- torch_tensor(e2, device = e1$device)
+  }
+  torch_matmul(e1, e2)
+}
+
 #' @export
 dim.torch_tensor <- function(x) {
   cpp_tensor_dim(x$ptr)

tests/testthat/test-operators.R

@@ -121,6 +121,22 @@ test_that("! works", {
   expect_equal_to_r(!y, c(TRUE, FALSE))
 })
 
+test_that("%*% works", {
+  if (getRversion() < "4.3.0") {
+    skip("%*% S3 operator for torch_tensor requires R >= 4.3.0")
+  }
+
+  lhs <- torch_tensor(matrix(c(1, 2, 3, 4), nrow = 2))
+  rhs_tensor <- torch_tensor(matrix(c(0, 1, 1, 0), nrow = 2))
+  rhs_matrix <- matrix(c(2, 0, 1, 2), nrow = 2)
+
+  lhs_r <- as.matrix(as_array(lhs))
+  rhs_tensor_r <- as.matrix(as_array(rhs_tensor))
+
+  expect_equal_to_r(lhs %*% rhs_tensor, lhs_r %*% rhs_tensor_r)
+  expect_equal_to_r(lhs %*% rhs_matrix, lhs_r %*% rhs_matrix)
+})
+
 test_that("dim works", {
   x <- torch_randn(c(2, 2))
   expect_equal(dim(x), c(2, 2))
@@ -357,6 +373,24 @@ test_that("| works", {
   expect_equal_to_r(x | 0, c(FALSE, TRUE))
 })
 
+test_that("%*% works on cuda", {
+  if (getRversion() < "4.3.0") {
+    skip("%*% S3 operator for torch_tensor requires R >= 4.3.0")
+  }
+
+  skip_if_cuda_not_available()
+
+  lhs <- torch_tensor(matrix(c(1, 2, 3, 4), nrow = 2), device = "cuda")
+  rhs_tensor <- torch_tensor(matrix(c(0, 1, 1, 0), nrow = 2), device = "cuda")
+  rhs_matrix <- matrix(c(2, 0, 1, 2), nrow = 2)
+
+  lhs_r <- as.matrix(as_array(lhs$cpu()))
+  rhs_tensor_r <- as.matrix(as_array(rhs_tensor$cpu()))
+
+  expect_equal_to_r(lhs %*% rhs_tensor, lhs_r %*% rhs_tensor_r)
+  expect_equal_to_r(lhs %*% rhs_matrix, lhs_r %*% rhs_matrix)
+})
+
 test_that("mean works", {
   x <- c(1, 2, 3, 4)