From 23e0d4a1a3557ebc21cb7b08f455bcecfea21944 Mon Sep 17 00:00:00 2001
From: Waldir Leoncio <waldir.leoncio@gmail.com>
Date: Wed, 18 Mar 2020 11:42:41 +0100
Subject: [PATCH] Improved handling of extra dimensions

---
 R/zeros_ones.R | 16 ++++++++++++++--
 man/ones.Rd    |  4 +++-
 man/zeros.Rd   |  4 +++-
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/R/zeros_ones.R b/R/zeros_ones.R
index 435cbe2..c4487bf 100644
--- a/R/zeros_ones.R
+++ b/R/zeros_ones.R
@@ -24,8 +24,14 @@ zeros_or_ones <- function(n, x) {
 #' the `zeros()` function on Matlab
 #' @param n1 number of rows
 #' @param n2 number of columns
+#' @param ... extra dimensions
 zeros <- function(n1, n2 = n1, ...) {
-    return(zeros_or_ones(c(n1, n2, ...), 0))
+    if (length(n1) == 1) {
+        n <- c(n1, n2, ...)
+    } else {
+        n <- n1
+    }
+    return(zeros_or_ones(n, 0))
 }
 
 #' @title Matrix of ones
@@ -33,6 +39,12 @@ zeros <- function(n1, n2 = n1, ...) {
 #' the `ones()` function on Matlab
 #' @param n1 number of rows
 #' @param n2 number of columns
+#' @param ... extra dimensions
 ones <- function(n1, n2 = n1, ...) {
-    return(zeros_or_ones(c(n1, n2, ...), 1))
+    if (length(n1) == 1) {
+        n <- c(n1, n2, ...)
+    } else {
+        n <- n1
+    }
+    return(zeros_or_ones(n, 1))
 }
\ No newline at end of file
diff --git a/man/ones.Rd b/man/ones.Rd
index e739460..b110cf8 100644
--- a/man/ones.Rd
+++ b/man/ones.Rd
@@ -4,12 +4,14 @@
 \alias{ones}
 \title{Matrix of ones}
 \usage{
-ones(n1, n2 = n1)
+ones(n1, n2 = n1, ...)
 }
 \arguments{
 \item{n1}{number of rows}
 
 \item{n2}{number of columns}
+
+\item{...}{extra dimensions}
 }
 \description{
 wrapper of `zeros_or_ones()` that replicates the behavior of 
diff --git a/man/zeros.Rd b/man/zeros.Rd
index 039b95b..a5fa748 100644
--- a/man/zeros.Rd
+++ b/man/zeros.Rd
@@ -4,12 +4,14 @@
 \alias{zeros}
 \title{Matrix of zeros}
 \usage{
-zeros(n1, n2 = n1)
+zeros(n1, n2 = n1, ...)
 }
 \arguments{
 \item{n1}{number of rows}
 
 \item{n2}{number of columns}
+
+\item{...}{extra dimensions}
 }
 \description{
 wrapper of `zeros_or_ones()` that replicates the behavior of