Improve handling of 3-dimensional output tensors + more explanation in error messages

Closes #7 (closed)

We can now use 3-dimensional tensors, without the need of expanding the tensor to have a 4th dimension 😇

For instance, this model now works:

# Input
x = tf.keras.Input(shape=[None, None, None], name="x")  # shape is [1, h, w, N]

# Compute norm on the last axis
y = tf.norm(x, axis=-1)  # shape is [1, h, w]

Before this MR, we had to expand y dimensions, or use keepdims=True in tf.norm in order to have y with shape [1, h, w, 1] ! not very intuitive...