Perform Einstein-style tensor operations

A unified interface for rearranging, reducing, and repeating tensor dimensions using Einstein notation-inspired expressions. This was directly adapted from the python einop package: https://github.com/cgarciae/einop

einop() auto detects the operation type based on the provided expression:

rearrange() when all input dimensions are present in output
reduce() when some input dimensions are missing from output
einops.repeat() when output contains dimensions not in input

Note that there are ongoing debates about the use of this function purely from the perspective of code readability and maintainability: https://github.com/arogozhnikov/einops/issues/84. Generally, some argue that the descriptive names of rearrange, reduce, and repeat encourage good practices, while others think that semantically einop() actually makes it clearer what the operation is doing, as opposed to mandating the use of these commonly used function names across many packages.

Usage

einop(
  x,
  expr,
  reduction = NULL,
  ...,
  .row_major = getOption("einops_row_major", FALSE)
)

Arguments

x: tensor: array, matrix, or list of arrays of the same shape and type
expr: string: reduction pattern
reduction: A string specifying the reduction operation (e.g., "mean", "sum", "max"). Required for reduce operations, ignored for rearrange and repeat operations.
...: either corresponding axes lengths or a single list of them.
.row_major: logical: whether to use row-major order for the output tensor. If TRUE, the operation is performed in row-major order, but the output will be in whatever order the parent framework uses (e.g. column-major for base::array()).

Value

A tensor with dimensions transformed according to the expression

Examples

if (requireNamespace("abind", quietly = TRUE)) {

# load a 3d tensor representing an rgb image
x <- get(data("einops_image"))[1, , , ]

# Rearrange dimensions
einop(x, "h w c -> c h w")

# Reduce dimensions
einop(x, "h w c -> h w", "mean")

# Repeat dimensions
einop(x[, , 1], "h w -> h w c", c = 3)

}