From b3309d6dfa6211e4a4d3dd533a0d27023028f225 Mon Sep 17 00:00:00 2001
From: Matt Streeter <mstreeter@google.com>
Date: Tue, 8 Aug 2023 14:45:34 -0700
Subject: [PATCH] Fix shape bug when raising multivariate enclosures to a
 power.

This fixes https://github.com/google/autobound/issues/10

Also bump version number in preparation for updating the pip package.
---
 autobound/__init__.py                  |  2 +-
 autobound/enclosure_arithmetic.py      |  3 ++-
 autobound/enclosure_arithmetic_test.py | 21 +++++++++++++++++++++
 3 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/autobound/__init__.py b/autobound/__init__.py
index 142cae7..a7ba5ab 100644
--- a/autobound/__init__.py
+++ b/autobound/__init__.py
@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "0.1.2"
+__version__ = "0.1.3"
diff --git a/autobound/enclosure_arithmetic.py b/autobound/enclosure_arithmetic.py
index a5e7894..100bcdb 100644
--- a/autobound/enclosure_arithmetic.py
+++ b/autobound/enclosure_arithmetic.py
@@ -328,7 +328,8 @@ def power(self, a: TaylorEnclosureLike, p: float) -> TaylorEnclosure:
       term_power_coefficient = functools.partial(
           _elementwise_term_power_coefficient, x_ndim=x_ndim,
           np_like=self.np_like)
-      multiplicative_identity = self.np_like.ones_like(self.trust_region[0])
+      a0 = a[0][0] if isinstance(a[0], tuple) else a[0]
+      multiplicative_identity = self.np_like.ones_like(a0)
       result = polynomials.integer_power(  # pytype: disable=wrong-arg-types
           a,
           p,
diff --git a/autobound/enclosure_arithmetic_test.py b/autobound/enclosure_arithmetic_test.py
index 88c5a53..d2d2466 100644
--- a/autobound/enclosure_arithmetic_test.py
+++ b/autobound/enclosure_arithmetic_test.py
@@ -225,6 +225,27 @@ def test_pairwise_batched_multiply(self, u, v, p, q, expected):
           0,
           (np.ones((3,)),)
       ),
+      (
+          100,
+          (np.zeros((2,)), np.ones((2,))),
+          (np.ones((3,)), np.ones((3, 2))),
+          0,
+          (np.ones((3,)),)
+      ),
+      (
+          100,
+          (np.zeros((2,)), np.ones((2,))),
+          (np.ones((3,)), np.ones((3, 2))),
+          1,
+          (np.ones((3,)), np.ones((3, 2))),
+      ),
+      (
+          100,
+          (np.zeros((2,)), np.ones((2,))),
+          (np.ones((3,)), np.ones((3, 2))),
+          2,
+          (np.ones((3,)), 2*np.ones((3, 2)), np.ones((3, 2, 2))),
+      ),
   )
   def test_power(self, max_degree, trust_region, enclosure, p, expected):
     for np_like in self.backends: