title: Tensorflow 2.0 Custom loss function with multiple inputs
body: <p>I am trying to optimize a model with the following two loss functions</p>
  
  <pre><code>def loss_1(pred, weights, logits):
      weighted_sparse_ce = kls.SparseCategoricalCrossentropy(from_logits=True)
      policy_loss = weighted_sparse_ce(pred, logits, sample_weight=advantages)
  </code></pre>
  
  <p>and</p>
  
  <pre><code>def loss_2(y_pred, y):
      return kls.mean_squared_error(y_pred, y)
  </code></pre>
  
  <p>however, because TensorFlow 2 expects loss function to be of the form</p>
  
  <pre><code>def fn(y_pred, y_true):
      ...
  </code></pre>
  
  <p>I am using a work-around for <code>loss_1</code> where I pack <code>pred</code> and <code>weights</code> into a single tensor before passing to <code>loss_1</code> in the call to <a href="https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/keras/Model#fit" rel="noreferrer"><code>model.fit</code></a> and then unpack them in <code>loss_1</code>. This is inelegant and nasty because <code>pred</code> and <code>weights</code> are of different data types and so this requires an additional cast, pack, un-pack and un-cast each time I call <a href="https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/keras/Model#fit" rel="noreferrer"><code>model.fit</code></a>.</p>
  
  <p>Furthermore, I am aware of the <code>sample_weight</code> argument to <code>fit</code>, which is kind of like the solution to <a href="https://stackoverflow.com/questions/57414114/implement-custom-loss-function-in-tensorflow-2-0">this question</a>. This might be a workable solution were it not for the fact that I am using two loss functions and I only want the <code>sample_weight</code> applied to one of them. Also, even if this were a solution, would it not be generalizable to other types of custom loss functions.</p>
  
  <hr>
  
  <p>All that being said, my question, said concisely, is:</p>
  
  <p>What is the best way to create a loss function with an arbitrary number of
  arguments in TensorFlow 2?</p>
  
  <p>Another thing I have tried is passing a <a href="https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/tuple" rel="noreferrer"><code>tf.tuple</code></a> but that also seems to violate TensorFlow's desires for a loss function input.</p>