def _log_matmul_exp(tensor_a, tensor_b):
  """Numerically stable equivalent of log-matmul-exp computation.

  Is equivalent to:
    tf.math.log(
      tf.matmul(tf.math.exp(tensor_a),
                tf.math.exp(tensor_b)))

  See:
  https://en.wikipedia.org/wiki/LogSumExp#log-sum-exp_trick_for_log-domain_calculations

  Args:
    tensor_a: <float>[L, M].
    tensor_b: <float>[M, N].

  Returns:
    <float>[L, N].
  """
  max_a = tf.math.reduce_max(tensor_a, axis=1, keepdims=True)
  max_b = tf.math.reduce_max(tensor_b, axis=0, keepdims=True)
  tensor_a -= max_a
  tensor_b -= max_b
  tensor_c = tf.matmul(tf.math.exp(tensor_a), tf.math.exp(tensor_b))
  tensor_c = tf.math.log(tensor_c)
  tensor_c += max_a
  tensor_c += max_b
  return tensor_c
Numerically stable equivalent of log-matmul-exp computation.

  Is equivalent to:
    tf.math.log(
      tf.matmul(tf.math.exp(tensor_a),
                tf.math.exp(tensor_b)))

  See:
  https://en.wikipedia.org/wiki/LogSumExp#log-sum-exp_trick_for_log-domain_calculations

  Args:
    tensor_a: <float>[L, M].
    tensor_b: <float>[M, N].

  Returns:
    <float>[L, N].
  
language/compgen/csl/model/weighted_model.py