From 8d99bc96c3c0fe5b14104ed261d1a7db4e78a592 Mon Sep 17 00:00:00 2001
From: "[zebinyang]" <yangzb2010@hku.hk>
Date: Tue, 19 Oct 2021 09:58:57 +0800
Subject: [PATCH] fix a bug in calculating clarity loss as sample weight is
 specified; update v0.5.6

---
 README.md           | 3 +++
 gaminet/__init__.py | 2 +-
 gaminet/gaminet.py  | 4 ++--
 setup.py            | 2 +-
 4 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 382f00f..aac2a74 100644
--- a/README.md
+++ b/README.md
@@ -7,6 +7,7 @@ The following environments are required:
 
 - Python 3.7 + (anaconda is preferable)
 - tensorflow>=2.0.0
+- tensorflow-lattice>=2.0.8
 - numpy>=1.15.2
 - pandas>=0.19.2
 - matplotlib>=3.1.3
@@ -16,6 +17,8 @@ The following environments are required:
 pip install gaminet
 ```
 
+To use it on GPU, conda install tensorflow==2.2, pip install tensorflow-lattice==2.0.8, conda install tensorflow-estimators==2.2
+
 ## Usage
 
 Import library
diff --git a/gaminet/__init__.py b/gaminet/__init__.py
index f7b102d..3fa7e37 100644
--- a/gaminet/__init__.py
+++ b/gaminet/__init__.py
@@ -2,5 +2,5 @@
 
 __all__ = ["GAMINet"]
 
-__version__ = '0.5.5'
+__version__ = '0.5.6'
 __author__ = 'Zebin Yang and Aijun Zhang'
diff --git a/gaminet/gaminet.py b/gaminet/gaminet.py
index 954ce67..12d95a1 100644
--- a/gaminet/gaminet.py
+++ b/gaminet/gaminet.py
@@ -159,8 +159,8 @@ def call(self, inputs, sample_weight=None, main_effect_training=False, interacti
                 a2 = tf.multiply(tf.gather(self.maineffect_outputs, [k2], axis=1), tf.gather(main_weights, [k2], axis=0))
                 b = tf.multiply(tf.gather(self.interact_outputs, [i], axis=1), tf.gather(interaction_weights, [i], axis=0))
                 if sample_weight is not None:
-                    self.clarity_loss += tf.abs(tf.reduce_mean(tf.multiply(tf.multiply(a1, b), sample_weight)))
-                    self.clarity_loss += tf.abs(tf.reduce_mean(tf.multiply(tf.multiply(a2, b), sample_weight)))
+                    self.clarity_loss += tf.abs(tf.reduce_mean(tf.multiply(tf.multiply(a1, b), sample_weight.reshape(-1, 1))))
+                    self.clarity_loss += tf.abs(tf.reduce_mean(tf.multiply(tf.multiply(a2, b), sample_weight.reshape(-1, 1))))
                 else:
                     self.clarity_loss += tf.abs(tf.reduce_mean(tf.multiply(a1, b)))
                     self.clarity_loss += tf.abs(tf.reduce_mean(tf.multiply(a2, b)))
diff --git a/setup.py b/setup.py
index 237cd95..7b3111b 100644
--- a/setup.py
+++ b/setup.py
@@ -14,7 +14,7 @@
 }
 
 setup(name='gaminet',
-      version='0.5.5',
+      version='0.5.6',
       description='Explainable Neural Networks based on Generalized Additive Models with Structured Interactions',
       url='https://github.com/ZebinYang/GAMINet',
       author='Zebin Yang',