Add test for tensorflow broadcasting

1 files changed, 4 insertions, 1 deletions

diff --git a/searchlib/src/test/files/integration/tensorflow/dropout/dropout.py b/searchlib/src/test/files/integration/tensorflow/dropout/dropout.py
index adbf29b9ab6..06ae4c4e5d5 100644
--- a/searchlib/src/test/files/integration/tensorflow/dropout/dropout.py
+++ b/searchlib/src/test/files/integration/tensorflow/dropout/dropout.py
@@ -16,8 +16,11 @@ X = tf.placeholder(tf.float32, shape=(None, n_inputs), name="X")
 y = tf.placeholder(tf.int64, shape=(None), name="y")
 training = tf.placeholder_with_default(False, shape=(), name='training')
 
+def leaky_relu_with_small_constant(z, name=None):
+    return tf.maximum(tf.constant(0.01, shape=[1]) * z, z, name=name)
+
 X_drop = tf.layers.dropout(X, dropout_rate, training=training, name="xdrop")
-output = tf.layers.dense(X_drop, n_outputs, name="outputs")
+output = tf.layers.dense(X_drop, n_outputs, activation=leaky_relu_with_small_constant, name="outputs")
 
 init = tf.global_variables_initializer()
 file_writer = tf.summary.FileWriter(logdir, tf.get_default_graph())