title: InvalidArgumentError: required broadcastable shapes [Op:Add] tensorflow model
body: <p>I am a using CLIP model. Where I have two models. One model output is <code>(20, 128, 256)</code> and the other one output is <code>(20, 256)</code>.</p>
  <pre><code>image_model_output = (20, 256)
  text_model_output = (20, 128, 256)
  </code></pre>
  <p>I use the following to calculate this</p>
  <pre><code>logits = (tf.matmul(caption_embeddings, image_embeddings, transpose_b=True))
  so it will be like `(20, 256) * (256, 128, 20)`
  it's ouput will be `(20, 128, 20)`
  </code></pre>
  <p>Similarly I calculate like this</p>
  <pre><code>images_similarity = tf.matmul(
          image_embeddings, image_embeddings, transpose_b=True
      )
  (Output)--&gt; (20, 256) * (256, 20) = (20,20) 
  </code></pre>
  <p>and this</p>
  <pre><code>captions_similarity = tf.matmul(
          caption_embeddings, caption_embeddings, transpose_b=True
      )
  (Output)--&gt; (20, 128, 256) * (256, 128, 20) = (20, 128, 128) 
  </code></pre>
  <p>The problem arises here</p>
  <pre><code>targets = keras.activations.softmax(
          (captions_similarity + images_similarity) / (2 * self.temperature)
      )
  </code></pre>
  <p>So do I need to change the activation function or there is any way to add these 3d matrices with different shapes?
  Sorry to technically explain like this but people with solid deep learning and machine learning backgorund will understand.</p>
  <h2>NOTE: After adding <code>axis 1</code> like this <code>tf.expand_dims(image_embeddings, axis=1)</code> the below part runs successfully</h2>
  <pre><code>targets = keras.activations.softmax(
      (captions_similarity + images_similarity) / (2 * self.temperature)
  )
  </code></pre>
  <p>However after this there is a loss funtion like below</p>
  <pre><code>captions_loss = keras.losses.categorical_crossentropy(
          y_true=targets, y_pred=logits, from_logits=True
      )
  </code></pre>
  <p>which generates this error</p>
  <pre><code>ValueError: Shapes (2, 128, 128) and (2, 128, 1) are incompatible
  </code></pre>
  <p>Is it possible to solve this error?</p>