Toolkit

library(rray)
library(magrittr)

Introduction

One of the big goals for rray is to be a general purpose toolkit for working with arrays. A requirement of this is that you shouldn’t have to opt into using rray objects to be able to take advantage of broadcasting or the ability to use any of the rray_*() functions. That requirement has been at the core of rray development, and it means that you can use base R vector/matrix/array objects with any of the rray functions, and still get a base R object back out.

x <- matrix(1:6, nrow = 2)
y <- matrix(1:2, nrow = 2)

# Base R matrices, added with broadcasting
rray_add(x, y)
#>      [,1] [,2] [,3]
#> [1,]    2    4    6
#> [2,]    4    6    8

Axes

Many of the functions in rray are applied “along an axis”. With base R, you might be used to the MARGIN argument when specifying the dimension to apply a function over. In rray, you’ll use axes (or axis, depending on the function). In short, these two are complements of one another. Ignoring the fact that rray doesn’t drop length 1 dimensions, notice that the values computed in the example below are the same, even though the dimensions to compute over look different.

x <- array(1:8, c(2, 2, 2))

rray_sum(x, axes = 1)
#> , , 1
#> 
#>      [,1] [,2]
#> [1,]    3    7
#> 
#> , , 2
#> 
#>      [,1] [,2]
#> [1,]   11   15

apply(x, MARGIN = c(2, 3), FUN = sum)
#>      [,1] [,2]
#> [1,]    3   11
#> [2,]    7   15

I find that axes is a more intuitive way to specify the dimensions. This is because with axes, you list the dimensions that you are allowing to change in some way. In the above example, axes = 1 was specified which guarantees that the result will have the same dimensions as x everywhere except in the first dimension, which will have length 1, no matter what. In other words, the dimensions change as: (2, 2, 2) -> (1, 2, 2).

Binding

Reducers aren’t the only functions with this axes guarantee. With rray_bind(), you specify the .axis that you want to bind along. This has the same guarantee that only the .axis specified will be changing. The only caveat here is that the inputs are first broadcast to a common dimension (ignoring the .axis dimension) before the binding is carried out. A few examples might be helpful:

# (5, 1)
x <- matrix(1:5)

# (3)
y <- 6:8

rray_bind(x, y, .axis = 1)
#>      [,1]
#> [1,]    1
#> [2,]    2
#> [3,]    3
#> [4,]    4
#> [5,]    5
#> [6,]    6
#> [7,]    7
#> [8,]    8

This works by first finding the common dimensions between x and y, ignoring the .axis dimension. So in this case the common dimension is (., 1) where the . represents whatever dimension is actually there for that object. The final result after binding the inputs together will also have a (., 1) shape, where the . will be the sum of all of the 1st dimension sizes of the inputs (in this case 5 + 3 = 8).

# (5, 1) where `. = 5` from x
x_broadcast <- rray_broadcast(x, c(5, 1))
x_broadcast
#>      [,1]
#> [1,]    1
#> [2,]    2
#> [3,]    3
#> [4,]    4
#> [5,]    5

# (3, 1) where `. = 3` from y
y_broadcast <- rray_broadcast(y, c(3, 1))
y_broadcast
#>      [,1]
#> [1,]    6
#> [2,]    7
#> [3,]    8

rray_bind(x_broadcast, y_broadcast, .axis = 1)
#>      [,1]
#> [1,]    1
#> [2,]    2
#> [3,]    3
#> [4,]    4
#> [5,]    5
#> [6,]    6
#> [7,]    7
#> [8,]    8

The .axis is actually taken account when finding the common dimensions. This means that you can bind along an axis that is higher in dimensionality than any of the inputs. For example, you can bind two matrices along the third dimension.

rray_bind(x, x, .axis = 3)
#> , , 1
#> 
#>      [,1]
#> [1,]    1
#> [2,]    2
#> [3,]    3
#> [4,]    4
#> [5,]    5
#> 
#> , , 2
#> 
#>      [,1]
#> [1,]    1
#> [2,]    2
#> [3,]    3
#> [4,]    4
#> [5,]    5

This works by finding the common dimension between x and x, which is just (5, 1), and then checking that against the .axis. If .axis is implicitly along a higher dimension, the common dimension is extended as needed. In this case it is extended to (5, 1, 1) to match the fact that you are binding along the third dimension. The inputs are broadcast to that shape, and then bound together.

x_broadcast <- rray_broadcast(x, c(5, 1, 1))
x_broadcast
#> , , 1
#> 
#>      [,1]
#> [1,]    1
#> [2,]    2
#> [3,]    3
#> [4,]    4
#> [5,]    5

rray_bind(x_broadcast, x_broadcast, .axis = 3)
#> , , 1
#> 
#>      [,1]
#> [1,]    1
#> [2,]    2
#> [3,]    3
#> [4,]    4
#> [5,]    5
#> 
#> , , 2
#> 
#>      [,1]
#> [1,]    1
#> [2,]    2
#> [3,]    3
#> [4,]    4
#> [5,]    5

Dimension dropping

One thing you will immediately notice when working with rray is that it often tries not to drop dimensions. This happens most apparently in subsetting, and in the reducers like rray_sum() and is in stark contrast to base R.

x <- matrix(1:6, ncol = 2)
x_rray <- as_rray(x)

x[1,]
#> [1] 1 4
x_rray[1,]
#> <rray<int>[,2][1]>
#>      [,1] [,2]
#> [1,]    1    4

apply(x, 2, sum)
#> [1]  6 15
rray_sum(x, axes = 1)
#>      [,1] [,2]
#> [1,]    6   15

The rationale for this has to do with how broadcasting works. When dimensions aren’t dropped, operations such as the following are natural:

x_rray / rray_sum(x_rray, axes = 1)
#> <rray<dbl>[,2][3]>
#>           [,1]      [,2]
#> [1,] 0.1666667 0.2666667
#> [2,] 0.3333333 0.3333333
#> [3,] 0.5000000 0.4000000

This doesn’t work as you might expect with base R, and can result in a tragic error since partial recycling kicks in and you don’t get an error.

x / apply(x, 2, sum)
#>           [,1]      [,2]
#> [1,] 0.1666667 0.2666667
#> [2,] 0.1333333 0.8333333
#> [3,] 0.5000000 0.4000000

# Equivalent to 
col_sums <- apply(x, 2, sum)
partially_recycled <- matrix(rep(col_sums, times = 3), ncol = 2)
partially_recycled
#>      [,1] [,2]
#> [1,]    6   15
#> [2,]   15    6
#> [3,]    6   15

x / partially_recycled
#>           [,1]      [,2]
#> [1,] 0.1666667 0.2666667
#> [2,] 0.1333333 0.8333333
#> [3,] 0.5000000 0.4000000

This works nicely in rray because of two reasons, both are necessary:

x_sum <- rray_sum(x_rray, axes = 1)

# When `/` is called, the inputs are broadcast to the same shape, like this
dim <- rray_dim_common(x_rray, x_sum)
x_broadcast <- rray_broadcast(x_rray, dim)
x_sum_broadcast <- rray_broadcast(x_sum, dim)

x_broadcast
#> <rray<int>[,2][3]>
#>      [,1] [,2]
#> [1,]    1    4
#> [2,]    2    5
#> [3,]    3    6

x_sum_broadcast
#> <rray<dbl>[,2][3]>
#>      [,1] [,2]
#> [1,]    6   15
#> [2,]    6   15
#> [3,]    6   15

x_broadcast / x_sum_broadcast
#> <rray<dbl>[,2][3]>
#>           [,1]      [,2]
#> [1,] 0.1666667 0.2666667
#> [2,] 0.3333333 0.3333333
#> [3,] 0.5000000 0.4000000

If you do want to drop dimensions with rray, you can explicitly call rray_squeeze() afterwards. As a rule of thumb, it is much easier to drop dimensions explicitly than it is to recover them.

x_rray %>%
  rray_sum(axes = 1) %>%
  rray_squeeze(axes = 1)
#> <rray<dbl>[2]>
#> [1]  6 15

If you come from NumPy, you might be used to reducers dropping the axis you reduce over. I think this is a mistake, and there have been a number of discussions on the NumPy forums about this choice. Here is why:

# This is the result you'd get in a NumPy sum. The 1st dimension is dropped
x_sum_dropped <- rray_squeeze(rray_sum(x_rray, axes = 1), axes = 1)
x_sum_dropped
#> <rray<dbl>[2]>
#> [1]  6 15

# Now broadcasting doesn't work!
x_rray / x_sum_dropped
#> Error: Non-broadcastable dimensions: (3, 2) and (2).

# So you have to add back the dimension
# (numpy has slightly cleaner ways to do this, but it's still an extra step)
x_sum_reshaped <- rray_expand(x_sum_dropped, 1)
x_rray / x_sum_reshaped
#> <rray<dbl>[,2][3]>
#>           [,1]      [,2]
#> [1,] 0.1666667 0.2666667
#> [2,] 0.3333333 0.3333333
#> [3,] 0.5000000 0.4000000

For the curious, Julia’s implementation of reducers works similarly to rray.