Merge pull request #320 from stan-dev/matrixOps

Overload matrix multiplication operator for rvars in R >= 4.3

DESCRIPTION

@@ -1,6 +1,6 @@
 Package: posterior
 Title: Tools for Working with Posterior Distributions
-Version: 1.5.0
+Version: 1.5.0.9000
 Date: 2023-10-31
 Authors@R: c(person("Paul-Christian", "Bürkner", email = "[email protected]", role = c("aut", "cre")),
              person("Jonah", "Gabry", email = "[email protected]", role = c("aut")),

NEWS.md

@@ -1,3 +1,11 @@
+# posterior (development version)
+
+### Enhancements
+
+* Matrix multiplication of `rvar`s can now be done with the base matrix 
+  multiplication operator (`%*%`) instead of `%**%` in R >= 4.3.
+
+
 # posterior 1.5.0
 
 ### Enhancements

R/rvar-math.R

@@ -135,8 +135,9 @@ Math.rvar_factor <- function(x, ...) {
 #' is used to efficiently multiply matrices across draws, so if either `x` or `y` is an [`rvar`],
 #' `x %**% y` will be much faster than `rdo(x %*% y)`.
 #'
-#' Because [`rvar`] is an S3 class and S3 classes cannot properly override `%*%`, [`rvar`]s use
-#' `%**%` for matrix multiplication.
+#' In R >= 4.3, you can also use `%*%` in place of `%**%` for matrix multiplication
+#' of [`rvar`]s. In R < 4.3, S3 classes cannot properly override `%*%`, so
+#' you must use `%**%` for matrix multiplication of [`rvar`]s.
 #'
 #' @return An [`rvar`] representing the matrix product of `x` and `y`.
 #'
@@ -208,6 +209,21 @@ Math.rvar_factor <- function(x, ...) {
   new_rvar(result, .nchains = nchains(x))
 }
 
+# This generic is not exported here as matrixOps is only in R >= 4.3, so we must
+# conditionally export it in .onLoad() for compatibility with older versions
+#' @rdname rvar-matmult
+#' @method matrixOps rvar
+matrixOps.rvar <- function(x, y) {
+  # as of R 4.3 this group generic is only used for %*%, but that may change
+  # in the future (crossprod and tcrossprod are planned), so we include this
+  # check for safety purposes
+  if (.Generic != "%*%") {
+    stop_no_call("Cannot apply `", .Generic, "` operator to rvar objects.")
+  }
+  x %**% y
+}
+
+
 #' Cholesky decomposition of random matrix
 #'
 #' Cholesky decomposition of an [`rvar`] containing a matrix.

R/zzz.R

@@ -9,4 +9,10 @@
   # See help("s3_register", package = "vctrs") for more information.
   vctrs::s3_register("dplyr::dplyr_reconstruct", "draws_df")
   vctrs::s3_register("ggplot2::scale_type", "rvar")
+
+  # S3 methods for matrixOps, which is a group generic that didn't exist
+  # until R 4.3, so we can't register it in NAMESPACE
+  if (getRversion() >= "4.3") {
+    registerS3method("matrixOps", "rvar", matrixOps.rvar, asNamespace("base"))
+  }
 }

tests/testthat/test-rvar-math.R

@@ -262,6 +262,15 @@ test_that("matrix multiplication works", {
 
 })
 
+test_that("%*% works in R >= 4.3", {
+  skip_if_not(getRversion() >= "4.3")
+
+  x <- rvar(array(1:24, dim = c(4,2,3)))
+  y <- rvar(array(c(2:13,12:1), dim = c(4,3,2)))
+
+  expect_equal(x %*% y, x %**% y)
+})
+
 test_that("diag works", {
   Sigma <- as_draws_rvars(example_draws("multi_normal"))$Sigma
 

vignettes/rvar.Rmd

@@ -207,8 +207,12 @@ mu + 1
 ```
 
 Matrix multiplication is also implemented (using a tensor product under the hood).
-Because the normal matrix multiplication operator in R (`%*%`) cannot be properly
-implemented for S3 datatypes, `rvar` uses `%**%` instead. A trivial example:
+In R < 4.3, the normal matrix multiplication operator (`%*%`) cannot be properly
+implemented for S3 datatypes, so `rvar` uses `%**%` instead. In R ≥ 4.3, which 
+does support matrix multiplication for S3 datatypes, you can use `%*%` to 
+matrix-multiply `rvar`s.
+
+A trivial example:
 
 ```{r matrix_mult}
 Sigma %**% diag(1:3)
@@ -223,7 +227,7 @@ includes:
 | Logical operators           | `&`, `|`, `!` |
 | Comparison operators        | `==`, `!=`, `<`, `<=`, `>=`, `>` |
 | Value matching              | `match()`, `%in%` |
-| Matrix multiplication       | `%**%` |
+| Matrix multiplication       | `%**%`, `%*%` (R ≥ 4.3 only) |
 | Basic functions             | `abs()`, `sign()`<br>`sqrt()`<br>`floor()`, `ceiling()`, `trunc()`, `round()`, `signif()` |
 | Logarithms and exponentials | `exp()`, `expm1()`<br>`log()`, `log10()`, `log2()`, `log1p()` |
 | Trigonometric functions     | `cos()`, `sin()`, `tan()`<br>`cospi()`, `sinpi()`, `tanpi()`<br>`acos()`, `asin()`, `atan()`|