Adding Support for New Ops

TensorFlow graphs are constructed by buildings ops that receive tensors as input and produce tensors as output. Internally, multiple kernels are registered for each op, which are implementations of the op for different architectures (e.g., CPU kernels, CUDA GPU kernels, etc.). TensorFlow Scala offers the Op.Builder interface to allow users to create arbitrary ops that the TensorFlow runtime supports.

For example, the implentation of tf.add(x, y) in TensorFlow Scala looks like this:

def add[T: TF : IsNotQuantized](
    x: Output[T],
    y: Output[T],
    name: String = "Add"
): Output[T] = {
  Op.Builder[(Output[T], Output[T]), Output[T]](
    opType = "Add",
    name = name,
    input = (x, y)
  ).setGradientFn(addGradient(_, _)(TF[T], IsNotQuantized[T]))
      .build().output
}

protected def addGradient[T: TF : IsNotQuantized](
    op: Op[(Output[T], Output[T]), Output[T]],
    outputGradient: Output[T]
): (Output[T], Output[T]) = {
  val xShape = tf.shape(op.input._1)
  val yShape = tf.shape(op.input._2)
  val (rx, ry) = tf.broadcastGradientArguments(xShape, yShape)
  (tf.reshape(tf.sum(outputGradient, rx), xShape),
      tf.reshape(tf.sum(outputGradient, ry), yShape))
}