title: What exactly is Keras's CategoricalCrossEntropy doing?
body: <p>I am porting a keras model over to <code>torch</code> and I'm having trouble replicating the exact behavior of keras/tensorflow's <code>'categorical_crossentropy'</code> after a softmax layer.  I have some workarounds for this problem, so I'm only interested in understanding what <em>exactly</em> tensorflow calculates when calculating categorical cross entropy.</p>
  <p>As a toy problem, I set up labels and predicted vectors</p>
  <pre class="lang-py prettyprint-override"><code>&gt;&gt;&gt; import tensorflow as tf
  &gt;&gt;&gt; from tensorflow.keras import backend as K
  &gt;&gt;&gt; import numpy as np
  
  
  &gt;&gt;&gt; true = np.array([[0.0, 1.0], [1.0, 0.0]])
  &gt;&gt;&gt; pred = np.array([[0.0, 1.0], [0.0, 1.0]])
  </code></pre>
  <p>And calculate the Categorical Cross Entropy with:</p>
  <pre class="lang-py prettyprint-override"><code>&gt;&gt;&gt; loss = tf.keras.losses.CategoricalCrossentropy()
  &gt;&gt;&gt; print(loss(pred, true).eval(session=K.get_session()))
  8.05904769897461
  </code></pre>
  <p>This differs from the analytical result</p>
  <pre class="lang-py prettyprint-override"><code>&gt;&gt;&gt; loss_analytical = -1*K.sum(true*K.log(pred))/pred.shape[0]
  &gt;&gt;&gt; print(loss_analytical.eval(session=K.get_session()))
  nan
  </code></pre>
  <p>I dug into the source code for keras/tf's cross entropy (see <a href="https://stackoverflow.com/questions/61558769/softmax-cross-entropy-implementation-in-tensorflow-github-source-code">Softmax Cross Entropy implementation in Tensorflow Github Source Code</a>) and found the c function at <a href="https://github.com/tensorflow/tensorflow/blob/c903b4607821a03c36c17b0befa2535c7dd0e066/tensorflow/compiler/tf2xla/kernels/softmax_op.cc" rel="noreferrer">https://github.com/tensorflow/tensorflow/blob/c903b4607821a03c36c17b0befa2535c7dd0e066/tensorflow/compiler/tf2xla/kernels/softmax_op.cc</a> line 116.  In that function, there is a comment:</p>
  <pre class="lang-c prettyprint-override"><code>// sum(-labels *
  // ((logits - max_logits) - log(sum(exp(logits - max_logits)))))
  // along classes
  // (The subtraction broadcasts along the batch dimension.)
  </code></pre>
  <p>And implementing that, I tried:</p>
  <pre class="lang-py prettyprint-override"><code>&gt;&gt;&gt; max_logits = K.max(pred, axis=0)
  &gt;&gt;&gt; max_logits = max_logits
  &gt;&gt;&gt; xent = K.sum(-true * ((pred - max_logits) - K.log(K.sum(K.exp(pred - max_logits)))))/pred.shape[0]
  
  &gt;&gt;&gt; print(xent.eval(session=K.get_session()))
  1.3862943611198906
  </code></pre>
  <p>I also tried to print the trace for <code>xent.eval(session=K.get_session())</code>, but the trace is ~95000  lines long.  So it begs the question: what exactly is keras/tf doing when calculating <code>'categorical_crossentropy'</code>? It makes sense that it doesn't return <code>nan</code>, that would cause training issues, but where does 8 come from?</p>