From d1959e7c9dd2bd91be735b952d5dd55cc04e278b Mon Sep 17 00:00:00 2001
From: Krishnamoorthi Jayakumar
 <43264985+JKrishnamoorthi@users.noreply.github.com>
Date: Sat, 20 Jan 2024 14:16:01 +0530
Subject: [PATCH] Re-shaped the mod_chi2's return to match the input shape
 (#760)

I'm merging this anyway//
---
 pisa/utils/stats.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/pisa/utils/stats.py b/pisa/utils/stats.py
index f4d0a4f9a..cba88805d 100644
--- a/pisa/utils/stats.py
+++ b/pisa/utils/stats.py
@@ -568,7 +568,7 @@ def mod_chi2(actual_values, expected_values):
         the inputs
 
     """
-
+    in_array_shape = np.shape(actual_values)
     actual_values = unp.nominal_values(actual_values).ravel()
     sigma = unp.std_devs(expected_values).ravel()
     expected_values = unp.nominal_values(expected_values).ravel()
@@ -587,7 +587,8 @@ def mod_chi2(actual_values, expected_values):
         m_chi2 = (
             (actual_values - expected_values)**2 / (sigma**2 + expected_values)
         )
-
+    
+    m_chi2 = m_chi2.reshape(in_array_shape)
     return m_chi2
 
 def correct_chi2(actual_values, expected_values):