heuristics to adapt dimensional relationships across labels, predictions, and sampleWeights

The Python implementation of TensorFlow has complex heuristics that try to Do What I Mean in aligning dimensional relationships across labels, predictions, and sampleWeights. How do we want to approach this in the Java implementation? I'll explain what I mean by "complex heuristics" in a minute, but first, here are some possible paths forward:

Status Quo
We try to mirror Python's logic exactly. Like Python, in the APIs where we apply heuristic dimensional adjustments, we don't document what dimensional relationships are permitted, nor how they are adjusted. Nor do we test them, beyond expected common cases.

Alternative: Mirror, Document, and Test
I see motivation to mirror Python's logic exactly. But I'd feel a lot better about this if we documented the behavior thoroughly in our APIs and methodically tested it.

Mirroring Python's logic would help support the following goals that we wrote into tensorflow-framework/README.md:

• If either you know how to implement a model in the Python Keras API, or you are reimplementing an existing Python Keras model in Java, you should be able to cleanly and naturally follow the same high-level structure in the framework API.

• Also, given some familiarity with patterns followed throughout the framework API, you should be able to easily translate every detail of a Python Keras implementation into the framework API.

Although I see the attraction of mirroring Python, I think its heuristics are so complex that they aren't practical to fully document or test. Also, speaking as someone new to deep learning but experienced in API design, the Do What I Mean style feels to me like poor design.

Alternative: Simplify and Provide Adapters
I'd feel a lot better about having most of our APIs support a consistent set of simple, commonly-needed alignments, and then providing explicit adapters to support less common variations.

What do I mean by "complex heuristics"?
Let's take Hinge as an example. It is a simple subclass of Loss.

We have added documentation of the permitted dimensional relationship across parameters (beyond what's in the Python function header), but we still only scratch the surface.

  /**
   * Generates an Operand that calculates the loss.
   *
 . . .
   *
   * @param labels the truth values or labels, must be either -1, 0, or 1. Values are expected to be
   *     -1 or 1. If binary (0 or 1) labels are provided they will be converted to -1 or 1.
   * @param predictions the predictions, values must be in the range [0. to 1.] inclusive.
   * @param sampleWeights Optional sampleWeights acts as a coefficient for the loss. If a scalar is
   *     provided, then the loss is simply scaled by the given value. If sampleWeights is a tensor
   *     of size [batch_size], then the total loss for each sample of the batch is rescaled by the
   *     corresponding element in the SampleWeights vector. If the shape of SampleWeights is
   *     [batch_size, d0, .. dN-1] (or can be broadcast to this shape), then each loss element of
   *     predictions is scaled by the corresponding value of SampleWeights. (Note on dN-1: all loss
   *     functions reduce by 1 dimension, usually axis=-1.)
   * @param <T> The data type of the predictions, sampleWeights and loss.
   * @return the loss
   * @throws IllegalArgumentException if the predictions are outside the range [0.-1.].
   */
  @Override
  public <T extends TNumber> Operand<T> call(
      Operand<? extends TNumber> labels, Operand<T> predictions, Operand<T> sampleWeights) {

So what dimensional relationships do we actually support amongst labels, predictions, and sampleWeights? The only way to find out is to trace through some fairly complex code. As is our typical pattern, this code mixes calls to LossesHelper.squeezeOrExpandDimensions with use of operations that support broadcasting.

The above call method delegates most of its work to Losses.hinge. That method calls squeezeOrExpandDimensions with just labels and predictions (omitting sampleWeights at this stage). The result is that, if labels and predictions differ in rank by 1, and if the higher-ranked of them has a final dimension of size 1, we squeeze that dimension. In any other case we simply continue, neither adjusting labels or predictions nor raising an exception.

The Losses.hinge method then carries out a tensor calculation that supports broadcasting:

    return tf.math.mean(
        tf.math.maximum(tf.math.sub(one, tf.math.mul(tLabels, predictions)), zero),
        tf.constant(-1));

This use of squeezeOrExpandDimensions would interact strangely with broadcasting in some edge cases. For example, imagine that predictions has shape (1) while labels has shape (2, 1). Then after squeezing the last dimension of labels, we'd broadcast predictions across the remaining dimension. But if predictions has shape (1) while labels has shape (2, 2), then we'd broadcast predictions across both dimensions of labels.

At any rate, after this tensor calculation, we return to Hinge.call, which does the following:

    Operand<T> losses = Losses.hinge(getTF(), tLabels, predictions);
    return LossesHelper.computeWeightedLoss(getTF(), losses, getReduction(), sampleWeights);

Digging into computeWeightedLoss, we find the following code:

    LossTuple<T> result = squeezeOrExpandDimensions(tf, null, loss, sampleWeight);
    loss = result.getTarget();
    sampleWeight = result.getSampleWeights();

    Operand<T> weightedLosses = tf.math.mul(loss, cast(tf, sampleWeight, inputType));
    loss = reduceWeightedLoss(tf, weightedLosses, reduction);

The input loss, above, is the output of an interaction between squeezeOrExpandDimensions and a broadcasting tensor operation. But now we apply the same pattern again! This time, we pass loss to squeezeOrExpandDimensions as though it were predictions. We don't pass any labels, but we do pass sampleWeight. This potentially adjusts sampleWeight as follows:

• If sampleWeight is not null and is not a scalar, then:
  • If the rank of sampleWeight is one greater than the rank of loss, then we try to squeeze the last dimension of sampleWeight. (Actually, right now we try to squeeze all of its dimensions, but that's different from Python and I have reported it as a bug.)
  • Else if the rank of sampleWeight is one less than the rank of loss, then we expand the last dimension of sampleWeight. (I.e. we append a dimension of size 1.)

After those adjustments, we do a multiplication that supports broadcasting:

    Operand<T> weightedLosses = tf.math.mul(loss, cast(tf, sampleWeight, inputType));

The full combination of these alternating squeezeOrExpandDimensions and broadcasting operations seems impractical to document or test. For example, imagine (although it's odd) that we originally had the following shapes:

• labels (1, 2)
• predictions (2, 1, 1)
• sampleWeights (2)

If I'm following the code correctly, here are the adjustments we'd make:

• Squeeze the last dimension of predictions, yielding shape (2, 1).
• Broadcast the first dimension of labels and the second of predictions, yielding a loss shape (2, 2).
• Expand the last dimension of sampleWeights, yielding (2, 1).
• Broadcast this new last dimension of sampleWeights across the loss.

Lies tensorflow-framework/src/main/java/org/tensorflow/framework/losses/Hinge.java, Losses.java und impl/LossesHelper.java, und prüfe anschließend tensorflow-framework/README.md. Verfolge, wie squeezeOrExpandDimensions und broadcasting bei labels, predictions und sampleWeights zusammenspielen; für den Abschluss ist ein festgelegter Ansatz mit entsprechendem API-Verhalten, entsprechender Dokumentation und Tests oder Adaptern erforderlich.