Functions such as matrix multiplication perform column x row operations and then do a reduction operation. I want to do something similar, but I would like to replace the multiplication and addition operators with something else, such as max and min. I have something that works but it seems "ugly" at best.
# Setup
a = tf.reshape(tf.range(0.0, 8.0), [4, 2])
b = tf.reshape(tf.range(4.0, 12.0), [2, 4])
# Baseline
tf.matmul(a, b)
<tf.Tensor: shape=(4, 4), dtype=float32, numpy=
array([[ 8., 9., 10., 11.],
[ 32., 37., 42., 47.],
[ 56., 65., 74., 83.],
[ 80., 93., 106., 119.]], dtype=float32)>
# Can this part be done better?
a_b = tf.reshape(a, [4, 1, 2])
b_b = tf.reshape(tf.transpose(b), [1, 4, 2])
# The result is at least correct
tf.reduce_sum(a_b * b_b, -1)
<tf.Tensor: shape=(4, 4), dtype=float32, numpy=
array([[ 8., 9., 10., 11.],
[ 32., 37., 42., 47.],
[ 56., 65., 74., 83.],
[ 80., 93., 106., 119.]], dtype=float32)>
# And it can be extended to be generic
tf.reduce_min(tf.maximum(a_b, b_b), -1)
<tf.Tensor: shape=(4, 4), dtype=float32, numpy=
array([[4., 5., 6., 7.],
[4., 5., 6., 7.],
[4., 5., 6., 7.],
[6., 6., 6., 7.]], dtype=float32)>
As shown above, I have a workable solution, but I would expect a framework like tensorflow to have a more generic method to do this or at least a way to produce the intermediate tensor. The tf.meshgrid function seems to "almost" do what I want but the arguments are limited to rank 1 tensors.
Additionally, the above solution does not scale well. Some profiling indicates that the intermediate tensors are materialized, even in graph mode.