diff --git a/tensorflow-core/tensorflow-core-api/src/bazel/api_def/api_def_SoftmaxCrossEntropyWithLogits.pbtxt b/tensorflow-core/tensorflow-core-api/src/bazel/api_def/api_def_SoftmaxCrossEntropyWithLogits.pbtxt
index 5dba2164cd6..e064562c0f2 100644
--- a/tensorflow-core/tensorflow-core-api/src/bazel/api_def/api_def_SoftmaxCrossEntropyWithLogits.pbtxt
+++ b/tensorflow-core/tensorflow-core-api/src/bazel/api_def/api_def_SoftmaxCrossEntropyWithLogits.pbtxt
@@ -1,6 +1,6 @@
 op {
   graph_op_name: "SoftmaxCrossEntropyWithLogits"
   endpoint {
-    name: "nn.raw.SoftmaxCrossEntropyWithLogits"
+    name: "nn.SoftmaxCrossEntropyWithLogits"
   }
 }
diff --git a/tensorflow-core/tensorflow-core-api/src/bazel/api_def/api_def_SparseSoftmaxCrossEntropyWithLogits.pbtxt b/tensorflow-core/tensorflow-core-api/src/bazel/api_def/api_def_SparseSoftmaxCrossEntropyWithLogits.pbtxt
index cf80ff77565..7627d5f6074 100644
--- a/tensorflow-core/tensorflow-core-api/src/bazel/api_def/api_def_SparseSoftmaxCrossEntropyWithLogits.pbtxt
+++ b/tensorflow-core/tensorflow-core-api/src/bazel/api_def/api_def_SparseSoftmaxCrossEntropyWithLogits.pbtxt
@@ -1,6 +1,6 @@
 op {
   graph_op_name: "SparseSoftmaxCrossEntropyWithLogits"
   endpoint {
-    name: "nn.raw.SparseSoftmaxCrossEntropyWithLogits"
+    name: "nn.SparseSoftmaxCrossEntropyWithLogits"
   }
 }
diff --git a/tensorflow-core/tensorflow-core-api/src/gen/annotations/org/tensorflow/op/NnOps.java b/tensorflow-core/tensorflow-core-api/src/gen/annotations/org/tensorflow/op/NnOps.java
index 8b25a15522f..cc4f53394d3 100644
--- a/tensorflow-core/tensorflow-core-api/src/gen/annotations/org/tensorflow/op/NnOps.java
+++ b/tensorflow-core/tensorflow-core-api/src/gen/annotations/org/tensorflow/op/NnOps.java
@@ -83,7 +83,6 @@
 import org.tensorflow.op.nn.Relu;
 import org.tensorflow.op.nn.Relu6;
 import org.tensorflow.op.nn.Selu;
-import org.tensorflow.op.nn.SigmoidCrossEntropyWithLogits;
 import org.tensorflow.op.nn.Softmax;
 import org.tensorflow.op.nn.SoftmaxCrossEntropyWithLogits;
 import org.tensorflow.op.nn.Softsign;
@@ -103,8 +102,6 @@
  * @see {@link Ops}
  */
 public final class NnOps {
-  public final NnRawOps raw;
-
   private final Scope scope;
 
   private final Ops ops;
@@ -112,7 +109,6 @@ public final class NnOps {
   NnOps(Ops ops) {
     this.scope = ops.scope();
     this.ops = ops;
-    raw = new NnRawOps(ops);
   }
 
   /**
@@ -1815,55 +1811,6 @@ public <T extends TNumber> Selu<T> selu(Operand<T> features) {
     return Selu.create(scope, features);
   }
 
-  /**
-   * Computes sigmoid cross entropy given <code>logits</code>.
-   *
-   *  <p>Measures the probability error in discrete classification tasks in which each class is
-   *  independent and not mutually exclusive. For instance, one could perform multilabel
-   *  classification where a picture can contain both an elephant and a dog at the same time.
-   *
-   *  <p>For brevity, let <code>x = logits</code>, <code>z = labels</code>. The logistic loss in
-   *  pseudo-code is
-   *
-   *  <pre>
-   *  z * -log(sigmoid(x)) + (1 - z) * -log(1 - sigmoid(x))
-   *   = z * -log(1 / (1 + exp(-x))) + (1 - z) * -log(exp(-x) / (1 + exp(-x)))
-   *   = z * log(1 + exp(-x)) + (1 - z) * (-log(exp(-x)) + log(1 + exp(-x)))
-   *   = z * log(1 + exp(-x)) + (1 - z) * (x + log(1 + exp(-x))
-   *   = (1 - z) * x + log(1 + exp(-x))
-   *   = x - x * z + log(1 + exp(-x))
-   *  </pre>
-   *
-   *  <p>For <code>x < 0</code>, to avoid overflow in <code>exp(-x)</code>, we reformulate the above
-   *
-   *  <pre>
-   *  x - x * z + log(1 + exp(-x))
-   *   = log(exp(x)) - x * z + log(1 + exp(-x))
-   *   = - x * z + log(1 + exp(x))
-   *  </pre>
-   *
-   *  <p>Hence, to ensure stability and avoid overflow, the implementation uses this equivalent
-   *  formulation
-   *
-   *  <pre>
-   *    max(x, 0) - x * z + log(1 + exp(-abs(x)))
-   *  </pre>
-   *
-   *  <p></ode>logits</code> and <code>labels</code> must have the same type and shape.
-   *
-   *  <p>
-   *
-   * @param labels the labels
-   * @param logits the logits of type float32 or float64
-   * @param <T> the type of labels and logits
-   * @return the component-wise logistic losses.
-   * @throws IllegalArgumentException if logits' and labels' do not have the same shape
-   */
-  public <T extends TNumber> Operand<T> sigmoidCrossEntropyWithLogits(Operand<T> labels,
-      Operand<T> logits) {
-    return SigmoidCrossEntropyWithLogits.sigmoidCrossEntropyWithLogits(scope, labels, logits);
-  }
-
   /**
    * Computes softmax activations.
    *  For each batch {@code i} and class {@code j} we have
@@ -1881,53 +1828,20 @@ public <T extends TNumber> Softmax<T> softmax(Operand<T> logits) {
   }
 
   /**
-   * Computes softmax cross entropy between <code>logits</code> and <code>labels</code>.
-   *
-   *  <p>Measures the probability error in discrete classification tasks in which the classes are
-   *  mutually exclusive (each entry is in exactly one class). For example, each CIFAR-10 image is
-   *  labeled with one and only one label: an image can be a dog or a truck, but not both.
-   *
-   *  <p><b>NOTE:</b>
-   *
-   *  <p>While the classes are mutually exclusive, their probabilities need not be. All that is
-   *  required is that each row of <code>labels</code> is a valid probability distribution. If they
-   *  are not, the computation of the gradient will be incorrect.
-   *
-   *  <p>If using exclusive <code>labels</code> (wherein one and only one class is true at a time),
-   *  see {@link org.tensorflow.op.NnOps#sparseSoftmaxCrossEntropyWithLogits}
-   *
-   *  <p>Usage:
-   *
-   *  <pre>
-   *    Operand&lt;TFloat32&gt; logits =
-   *        tf.constant(new float[][] {{4.0F, 2.0F, 1.0F}, {0.0F, 5.0F, 1.0F}} );
-   *    Operand&lt;TFloat32&gt; labels =
-   *        tf.constant(new float[][] {{1.0F, 0.0F, 0.0F}, {0.0F, 0.8F, 0.2F}} );
-   *    Operand&lt;TFloat32&gt; output =
-   *        tf.nn.softmaxCrossEntropyWithLogits(labels, logits, -1);
-   *    // output Shape = [2]
-   *    // dataType = FLOAT (1)
-   *    // values { 0.169846, 0.824745 }
-   *  </pre>
-   *
-   *  <p>Backpropagation will happen into both <code>logits</code> and <code>labels</code>. To
-   *  disallow backpropagation into <code>labels</code>, pass label tensors through <code>
-   *  tf.stopGradient</code> before feeding it to this function.
+   * Computes softmax cross entropy cost and gradients to backpropagate.
+   *  Inputs are the logits, not probabilities.
    *
-   * @param labels Each vector along the class dimension should hold a valid probability
-   *      distribution e.g. for the case in which labels are of shape <code>[batch_size, num_classes]
-   *      </code>, each row of <code>labels[i]</code> must be a valid probability distribution.
-   * @param logits Per-label activations, typically a linear output. These activation energies are
-   *      interpreted as unnormalized log probabilities.
-   * @param axis The class dimension. -1 is the last dimension.
-   * @param <T> the number type of the operands
-   * @return the softmax cross entropy loss. Its type is the same as <code>logits</code> and its
-   *      shape is the same as <code>labels</code> except that it does not have the last dimension of
-   *      <code>labels</code>.
+   * @param <T> data type for {@code loss} output
+   * @param features batch_size x num_classes matrix
+   * @param labels batch_size x num_classes matrix
+   *  The caller must ensure that each batch of labels represents a valid
+   *  probability distribution.
+   * @param <T> data type for {@code SoftmaxCrossEntropyWithLogits} output and operands
+   * @return a new instance of SoftmaxCrossEntropyWithLogits
    */
-  public <T extends TNumber, U extends TNumber> Operand<T> softmaxCrossEntropyWithLogits(
-      Operand<U> labels, Operand<T> logits, int axis) {
-    return SoftmaxCrossEntropyWithLogits.softmaxCrossEntropyWithLogits(scope, labels, logits, axis);
+  public <T extends TNumber> SoftmaxCrossEntropyWithLogits<T> softmaxCrossEntropyWithLogits(
+      Operand<T> features, Operand<T> labels) {
+    return SoftmaxCrossEntropyWithLogits.create(scope, features, labels);
   }
 
   /**
@@ -2114,50 +2028,23 @@ public <T extends TType> SpaceToDepth<T> spaceToDepth(Operand<T> input, Long blo
   }
 
   /**
-   * Computes sparse softmax cross entropy between <code>logits</code> and <code>labels</code>.
-   *
-   *  <p>Measures the probability error in discrete classification tasks in which the classes are
-   *  mutually exclusive (each entry is in exactly one class). For example, each CIFAR-10 image is
-   *  labeled with one and only one label: an image can be a dog or a truck, but not both.
-   *
-   *  <p><b>NOTE:</b>
-   *
-   *  <p>For this operation, the probability of a given label is considered exclusive. That is, soft
-   *  classes are not allowed, and the <code>labels</code> vector must provide a single specific
-   *  index for the true class for each row of <code>logits</code> (each minibatch entry). For soft
-   *  softmax classification with a probability distribution for each entry, {@link
-   *  org.tensorflow.op.NnOps#softmaxCrossEntropyWithLogits}.
-   *
-   *  <p><b>WARNING:</b>
+   * Computes softmax cross entropy cost and gradients to backpropagate.
+   *  Unlike {@code SoftmaxCrossEntropyWithLogits}, this operation does not accept
+   *  a matrix of label probabilities, but rather a single label per row
+   *  of features.  This label is considered to have probability 1.0 for the
+   *  given row.
+   *  <p>Inputs are the logits, not probabilities.
    *
-   *  <p>This op expects unscaled logits, since it performs a <code>softmax</code> on <code>logits
-   *  </code> internally for efficiency. Do not call this op with the output of <code>softmax</code>,
-   *  as it will produce incorrect results.
-   *
-   *  <p>A common use case is to have logits of shape <code>[batchSize, numClasses]</code> and have
-   *  labels of shape <code>[batchSize]</code>, but higher dimensions are supported, in which case
-   *  the <code>dim</code>-th dimension is assumed to be of size <code>numClasses</code>. <code>
-   *  logits</code> must have the <cod>dataType</cod> of <code>TFloat16</code>, <code>TFloat32</code>
-   *  , or <code>TFloat64</code>, and <code>labels</code> must have the dtype of <code>TInt32</code>
-   *  or <code>TInt64</code>.
-   *
-   * @param labels <code>Tensor</code> of shape <code>[d_0, d_1, ..., d_{r-1}]</code> (where <code>r
-   *      </code> is rank of <code>labels</code> and result) and the dataType is <code>TInt32</code>
-   *      or <code>TInt64</code>. Each entry in <code>labels</code> must be an index in <code>[0,
-   *      numClasses)</code>. Other values will raise an exception when this op is run on CPU, and
-   *      return <code>NaN</code> for corresponding loss and gradient rows on GPU.
-   * @param logits Per-label activations (typically a linear output) of shape <code>[d_0, d_1, ...,
-   *      d_{r-1}, numClasses]</code> and dataType of <code>TFloat16</code>, <code>TFloat32</code>,
-   *      or <code>TFloat64</code>. These activation energies are interpreted as unnormalized log
-   *      probabilities.
-   * @return A <code>Tensor</code> of the same shape as <code>labels</code> and of the same type as
-   *      <code>logits</code> with the softmax cross entropy loss.
-   * @throws IllegalArgumentException If logits are scalars (need to have rank >= 1) or if the rank
-   *      of the labels is not equal to the rank of the logits minus one.
-   */
-  public <T extends TNumber, U extends TNumber> Operand sparseSoftmaxCrossEntropyWithLogits(
-      Operand<T> labels, Operand<U> logits) {
-    return SparseSoftmaxCrossEntropyWithLogits.sparseSoftmaxCrossEntropyWithLogits(scope, labels, logits);
+   * @param <T> data type for {@code loss} output
+   * @param features batch_size x num_classes matrix
+   * @param labels batch_size vector with values in [0, num_classes).
+   *  This is the label for the given minibatch entry.
+   * @param <T> data type for {@code SparseSoftmaxCrossEntropyWithLogits} output and operands
+   * @return a new instance of SparseSoftmaxCrossEntropyWithLogits
+   */
+  public <T extends TNumber> SparseSoftmaxCrossEntropyWithLogits<T> sparseSoftmaxCrossEntropyWithLogits(
+      Operand<T> features, Operand<? extends TNumber> labels) {
+    return SparseSoftmaxCrossEntropyWithLogits.create(scope, features, labels);
   }
 
   /**
diff --git a/tensorflow-core/tensorflow-core-api/src/gen/annotations/org/tensorflow/op/NnRawOps.java b/tensorflow-core/tensorflow-core-api/src/gen/annotations/org/tensorflow/op/NnRawOps.java
deleted file mode 100644
index c287459c460..00000000000
--- a/tensorflow-core/tensorflow-core-api/src/gen/annotations/org/tensorflow/op/NnRawOps.java
+++ /dev/null
@@ -1,83 +0,0 @@
-// Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-// ==============================================================================
-//
-// This class has been generated, DO NOT EDIT!
-//
-package org.tensorflow.op;
-
-import org.tensorflow.Operand;
-import org.tensorflow.op.nn.raw.SoftmaxCrossEntropyWithLogits;
-import org.tensorflow.op.nn.raw.SparseSoftmaxCrossEntropyWithLogits;
-import org.tensorflow.types.family.TNumber;
-
-/**
- * An API for building {@code nn.raw} operations as {@link Op Op}s
- *
- * @see {@link Ops}
- */
-public final class NnRawOps {
-  private final Scope scope;
-
-  private final Ops ops;
-
-  NnRawOps(Ops ops) {
-    this.scope = ops.scope();
-    this.ops = ops;
-  }
-
-  /**
-   * Computes softmax cross entropy cost and gradients to backpropagate.
-   *  Inputs are the logits, not probabilities.
-   *
-   * @param <T> data type for {@code loss} output
-   * @param features batch_size x num_classes matrix
-   * @param labels batch_size x num_classes matrix
-   *  The caller must ensure that each batch of labels represents a valid
-   *  probability distribution.
-   * @param <T> data type for {@code SoftmaxCrossEntropyWithLogits} output and operands
-   * @return a new instance of SoftmaxCrossEntropyWithLogits
-   */
-  public <T extends TNumber> SoftmaxCrossEntropyWithLogits<T> softmaxCrossEntropyWithLogits(
-      Operand<T> features, Operand<T> labels) {
-    return SoftmaxCrossEntropyWithLogits.create(scope, features, labels);
-  }
-
-  /**
-   * Computes softmax cross entropy cost and gradients to backpropagate.
-   *  Unlike {@code SoftmaxCrossEntropyWithLogits}, this operation does not accept
-   *  a matrix of label probabilities, but rather a single label per row
-   *  of features.  This label is considered to have probability 1.0 for the
-   *  given row.
-   *  <p>Inputs are the logits, not probabilities.
-   *
-   * @param <T> data type for {@code loss} output
-   * @param features batch_size x num_classes matrix
-   * @param labels batch_size vector with values in [0, num_classes).
-   *  This is the label for the given minibatch entry.
-   * @param <T> data type for {@code SparseSoftmaxCrossEntropyWithLogits} output and operands
-   * @return a new instance of SparseSoftmaxCrossEntropyWithLogits
-   */
-  public <T extends TNumber> SparseSoftmaxCrossEntropyWithLogits<T> sparseSoftmaxCrossEntropyWithLogits(
-      Operand<T> features, Operand<? extends TNumber> labels) {
-    return SparseSoftmaxCrossEntropyWithLogits.create(scope, features, labels);
-  }
-
-  /**
-   * Get the parent {@link Ops} object.
-   */
-  public final Ops ops() {
-    return ops;
-  }
-}
diff --git a/tensorflow-core/tensorflow-core-api/src/gen/annotations/org/tensorflow/op/Ops.java b/tensorflow-core/tensorflow-core-api/src/gen/annotations/org/tensorflow/op/Ops.java
index a4a7f5d6dbc..fde138e7296 100644
--- a/tensorflow-core/tensorflow-core-api/src/gen/annotations/org/tensorflow/op/Ops.java
+++ b/tensorflow-core/tensorflow-core-api/src/gen/annotations/org/tensorflow/op/Ops.java
@@ -367,10 +367,10 @@ public final class Ops {
 
   public final SignalOps signal;
 
-  public final TrainOps train;
-
   public final QuantizationOps quantization;
 
+  public final TrainOps train;
+
   private final Scope scope;
 
   private Ops(Scope scope) {
@@ -393,8 +393,8 @@ private Ops(Scope scope) {
     math = new MathOps(this);
     audio = new AudioOps(this);
     signal = new SignalOps(this);
-    train = new TrainOps(this);
     quantization = new QuantizationOps(this);
+    train = new TrainOps(this);
   }
 
   /**
diff --git a/tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/internal/c_api/global/tensorflow.java b/tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/internal/c_api/global/tensorflow.java
index 2441bc1af65..b345ab4dad2 100644
--- a/tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/internal/c_api/global/tensorflow.java
+++ b/tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/internal/c_api/global/tensorflow.java
@@ -3150,7 +3150,7 @@ public static native void TF_RegisterFilesystemPlugin(
 // TF_InitKernel to do op/kernel registration.
 // Plugin should implement TF_InitKernel to register kernels. This function
 // should register all kernels in a plugin.
-public static native void TF_InitKernel();
+
 // Targeting ../Create_func_TF_OpKernelConstruction.java
 
 
diff --git a/tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/op/nn/raw/SoftmaxCrossEntropyWithLogits.java b/tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/op/nn/SoftmaxCrossEntropyWithLogits.java
similarity index 96%
rename from tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/op/nn/raw/SoftmaxCrossEntropyWithLogits.java
rename to tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/op/nn/SoftmaxCrossEntropyWithLogits.java
index 0f2d50289b9..8c3ef2da29a 100644
--- a/tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/op/nn/raw/SoftmaxCrossEntropyWithLogits.java
+++ b/tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/op/nn/SoftmaxCrossEntropyWithLogits.java
@@ -15,7 +15,7 @@
 
 // This class has been generated, DO NOT EDIT!
 
-package org.tensorflow.op.nn.raw;
+package org.tensorflow.op.nn;
 
 import org.tensorflow.Operand;
 import org.tensorflow.Operation;
@@ -34,7 +34,7 @@
  * @param <T> data type for {@code loss} output
  */
 @Operator(
-    group = "nn.raw"
+    group = "nn"
 )
 public final class SoftmaxCrossEntropyWithLogits<T extends TNumber> extends RawOp {
   /**
@@ -50,7 +50,7 @@ private SoftmaxCrossEntropyWithLogits(Operation operation) {
     super(operation);
     int outputIdx = 0;
     loss = operation.output(outputIdx++);
-    backprop = operation.output(outputIdx++);
+    backprop = operation.output(outputIdx);
   }
 
   /**
diff --git a/tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/op/nn/raw/SparseSoftmaxCrossEntropyWithLogits.java b/tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/op/nn/SparseSoftmaxCrossEntropyWithLogits.java
similarity index 97%
rename from tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/op/nn/raw/SparseSoftmaxCrossEntropyWithLogits.java
rename to tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/op/nn/SparseSoftmaxCrossEntropyWithLogits.java
index 30eec4c3c05..84feae55726 100644
--- a/tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/op/nn/raw/SparseSoftmaxCrossEntropyWithLogits.java
+++ b/tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/op/nn/SparseSoftmaxCrossEntropyWithLogits.java
@@ -15,7 +15,7 @@
 
 // This class has been generated, DO NOT EDIT!
 
-package org.tensorflow.op.nn.raw;
+package org.tensorflow.op.nn;
 
 import org.tensorflow.Operand;
 import org.tensorflow.Operation;
@@ -38,7 +38,7 @@
  * @param <T> data type for {@code loss} output
  */
 @Operator(
-    group = "nn.raw"
+    group = "nn"
 )
 public final class SparseSoftmaxCrossEntropyWithLogits<T extends TNumber> extends RawOp {
   /**
@@ -54,7 +54,7 @@ private SparseSoftmaxCrossEntropyWithLogits(Operation operation) {
     super(operation);
     int outputIdx = 0;
     loss = operation.output(outputIdx++);
-    backprop = operation.output(outputIdx++);
+    backprop = operation.output(outputIdx);
   }
 
   /**
diff --git a/tensorflow-core/tensorflow-core-api/src/gen/resources/ops.pb b/tensorflow-core/tensorflow-core-api/src/gen/resources/ops.pb
index 4c3e6bef038..cc39a35b203 100644
Binary files a/tensorflow-core/tensorflow-core-api/src/gen/resources/ops.pb and b/tensorflow-core/tensorflow-core-api/src/gen/resources/ops.pb differ
diff --git a/tensorflow-framework/src/main/java/org/tensorflow/framework/losses/Losses.java b/tensorflow-framework/src/main/java/org/tensorflow/framework/losses/Losses.java
index 9aa94cf7fcf..d23059b88fd 100644
--- a/tensorflow-framework/src/main/java/org/tensorflow/framework/losses/Losses.java
+++ b/tensorflow-framework/src/main/java/org/tensorflow/framework/losses/Losses.java
@@ -14,9 +14,12 @@
 =======================================================================*/
 package org.tensorflow.framework.losses;
 
+import static org.tensorflow.framework.utils.CastHelper.cast;
+
 import org.tensorflow.Operand;
 import org.tensorflow.framework.losses.impl.LossTuple;
 import org.tensorflow.framework.losses.impl.LossesHelper;
+import org.tensorflow.framework.op.FrameworkOps;
 import org.tensorflow.ndarray.Shape;
 import org.tensorflow.op.Ops;
 import org.tensorflow.op.core.ReduceAll;
@@ -28,8 +31,6 @@
 import org.tensorflow.types.TInt64;
 import org.tensorflow.types.family.TNumber;
 
-import static org.tensorflow.framework.utils.CastHelper.cast;
-
 /** Built-in loss functions. */
 public class Losses {
 
@@ -181,7 +182,10 @@ public static <T extends TNumber> Operand<T> binaryCrossentropy(
    */
   private static <T extends TNumber> Operand<T> binaryCrossentropyHelper(
       Ops tf, Operand<T> target, Operand<T> output, boolean fromLogits) {
-    if (fromLogits) return tf.nn.sigmoidCrossEntropyWithLogits(target, output);
+    FrameworkOps ftf = FrameworkOps.create(tf);
+    if (fromLogits) {
+      return ftf.nn.sigmoidCrossEntropyWithLogits(target, output);
+    }
 
     /* TODO - skip this logic for now. It requires walking back the inputs which is not yet possible
     if (!(output instanceof Variable) && (!tf.scope().env().isEager())) {
@@ -235,6 +239,7 @@ public static <T extends TNumber> Operand<T> categoricalCrossentropy(
       boolean fromLogits,
       float labelSmoothing,
       int axis) {
+    FrameworkOps ftf = FrameworkOps.create(tf);
     Class<T> predictionType = predictions.type();
     Operand<T> tLabels = cast(tf, labels, predictionType);
     LossTuple<T> ops = LossesHelper.squeezeOrExpandDimensions(tf, tLabels, predictions, null);
@@ -245,7 +250,7 @@ public static <T extends TNumber> Operand<T> categoricalCrossentropy(
       tLabels = smoothCategoricalLabels(tf, tLabels, labelSmoothing);
     }
     if (fromLogits) {
-      return tf.nn.softmaxCrossEntropyWithLogits(tLabels, predictions, axis);
+      return ftf.nn.softmaxCrossEntropyWithLogits(tLabels, predictions, axis);
     }
     /* TODO
     if (!(predictions instanceof Variable) && (!tf.scope().env().isEager())) {
@@ -515,6 +520,7 @@ public static <T extends TNumber> Operand<T> sparseCategoricalCrossentropy(
       Operand<T> predictions,
       boolean fromLogits,
       int axis) {
+    FrameworkOps ftf = FrameworkOps.create(tf);
     Class<T> predictionType = predictions.type();
     Operand<T> epsilonConst = cast(tf, tf.constant(EPSILON), predictionType);
     Operand<T> one = cast(tf, tf.constant(1), predictionType);
@@ -569,8 +575,7 @@ public static <T extends TNumber> Operand<T> sparseCategoricalCrossentropy(
                   new long[] {-1L, predictionsShape.size(predictionsShape.numDimensions() - 1)}));
     }
 
-    @SuppressWarnings("unchecked")
-    Operand<T> loss = tf.nn.sparseSoftmaxCrossEntropyWithLogits(iLabels, predictions);
+    Operand<T> loss = ftf.nn.sparseSoftmaxCrossEntropyWithLogits(iLabels, predictions);
     if (updateShape && predictionsRank >= 3) {
       Shape newShape = predictionsShape.take(predictionsShape.numDimensions() - 1);
       loss = tf.reshape(loss, tf.constant(newShape));
diff --git a/tensorflow-framework/src/main/java/org/tensorflow/framework/op/FrameworkOps.java b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/FrameworkOps.java
new file mode 100644
index 00000000000..7b6322d0f0d
--- /dev/null
+++ b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/FrameworkOps.java
@@ -0,0 +1,149 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+=======================================================================*/
+package org.tensorflow.framework.op;
+
+import org.tensorflow.DeviceSpec;
+import org.tensorflow.EagerSession;
+import org.tensorflow.ExecutionEnvironment;
+import org.tensorflow.op.Op;
+import org.tensorflow.op.Ops;
+import org.tensorflow.op.Scope;
+
+/**
+ * An API for building framework operations as {@link Op Op}s
+ *
+ * <p>These are higher level ops that may invoke core ops. Higher level Ops may perform the
+ * operation solely in the TensorFlow framework or do preprocessing of the Operands before invoking
+ * a core level Op.
+ */
+public class FrameworkOps {
+  public final Ops core;
+  public final NnOps nn;
+  public final SetOps sets;
+  public final MathOps math;
+  public final LinalgOps linalg;
+  private final Scope scope;
+
+  /**
+   * Creates a FrameworkOps instance with the provided scope
+   *
+   * @param scope the scope
+   */
+  private FrameworkOps(Scope scope) {
+    this.core = Ops.create(scope.env());
+    this.scope = scope;
+    nn = new NnOps(this);
+    sets = new SetOps(this);
+    math = new MathOps(this);
+    linalg = new LinalgOps(this);
+  }
+
+  /**
+   * Creates a FrameworkOps instance based on the provided Core Ops
+   *
+   * @param core The TensorFlow Core Ops
+   */
+  private FrameworkOps(Ops core) {
+    this.core = core;
+    this.scope = core.scope();
+    nn = new NnOps(this);
+    sets = new SetOps(this);
+    math = new MathOps(this);
+    linalg = new LinalgOps(this);
+  }
+
+  /**
+   * Creates an API for building operations in the provided execution environment
+   *
+   * @param env the exection environment
+   * @return the FrameworkOps
+   */
+  public static FrameworkOps create(ExecutionEnvironment env) {
+    return new FrameworkOps(new Scope(env));
+  }
+
+  /**
+   * Creates an API for building operations in the default eager execution environment
+   *
+   * <p>Invoking this method is equivalent to {@code
+   * FrameworkOps.create(EagerSession.getDefault())}.
+   *
+   * @return the FrameworkOps
+   */
+  public static FrameworkOps create() {
+    return new FrameworkOps(new Scope(EagerSession.getDefault()));
+  }
+
+  /**
+   * Creates an API for building operations in the default eager execution environment
+   *
+   * @param coreOps the TensorFlow core Ops
+   * @return the FrameworkOps
+   */
+  public static FrameworkOps create(Ops coreOps) {
+    return new FrameworkOps(coreOps);
+  }
+
+  /**
+   * Returns the current {@link Scope scope} of this API
+   *
+   * @return the current {@link Scope scope} of this API
+   */
+  public final Scope scope() {
+    return scope;
+  }
+
+  /**
+   * Returns an API that builds operations with the provided name prefix.
+   *
+   * <p>@link Scope#withSubScope(String)}
+   */
+  public FrameworkOps withSubScope(String childScopeName) {
+    return new FrameworkOps(scope.withSubScope(childScopeName));
+  }
+
+  /**
+   * Returns an API that uses the provided name for an op.
+   *
+   * <p>{@link Scope#withName(String)}
+   */
+  public FrameworkOps withName(String opName) {
+    return new FrameworkOps(scope.withName(opName));
+  }
+
+  /**
+   * Returns an API that places the created operations on the device(s) matching the provided spec.
+   *
+   * <p>{@link Scope#withDevice(DeviceSpec)}
+   *
+   * @param deviceSpec the device specification for the scope
+   * @return the FrameworkOps
+   */
+  public FrameworkOps withDevice(DeviceSpec deviceSpec) {
+    return new FrameworkOps(scope.withDevice(deviceSpec));
+  }
+
+  /**
+   * Returns an API that adds operations to the graph with the provided control dependencies.
+   *
+   * <p>{@link Scope#withControlDependencies(Iterable)}
+   *
+   * @param controls the operations
+   * @return the FrameworkOps
+   */
+  public FrameworkOps withControlDependencies(Iterable<Op> controls) {
+    return new FrameworkOps(scope.withControlDependencies(controls));
+  }
+}
diff --git a/tensorflow-framework/src/main/java/org/tensorflow/framework/op/LinalgOps.java b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/LinalgOps.java
new file mode 100644
index 00000000000..bc5118e494e
--- /dev/null
+++ b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/LinalgOps.java
@@ -0,0 +1,206 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+=======================================================================*/
+package org.tensorflow.framework.op;
+
+import org.tensorflow.Operand;
+import org.tensorflow.framework.op.linalg.MatMul;
+import org.tensorflow.framework.utils.SparseTensor;
+import org.tensorflow.op.Scope;
+import org.tensorflow.types.family.TNumber;
+
+public class LinalgOps {
+  private final Scope scope;
+
+  private final FrameworkOps frameworkOps;
+
+  /**
+   * Creates Framework linear algegra Operations
+   *
+   * @param frameworkOps the TensorFLow framework Ops
+   */
+  LinalgOps(FrameworkOps frameworkOps) {
+    this.scope = frameworkOps.scope();
+    this.frameworkOps = frameworkOps;
+  }
+
+  /**
+   * Multiplies matrix {@code a} by matrix {@code b}, producing {@code a} * {@code b }.
+   *
+   * <p>The inputs must, following any transpositions, be tensors of {@code rank >= 2} where the
+   * inner 2 dimensions specify valid matrix multiplication dimensions, and any further outer
+   * dimensions specify matching batch size.
+   *
+   * <p>Both matrices must be of the same type. The supported types are: {@code TFloat16}, {@code
+   * TFloat32}, {@code TFloat64}, {@code TInt32}.
+   *
+   * <p>Either matrix can be transposed or adjointed (conjugated and transposed) on the fly by
+   * setting one of the corresponding flag to true. These are false by default.
+   *
+   * <p>A simple 2-D tensor matrix multiplication:
+   *
+   * <pre>{@code
+   * Operand<TFloat64> a = tf.constant(new double[][] {
+   *         {-8.944851},
+   *         {4.1711287},
+   *         {-0.22380222}
+   *     });
+   * Operand<TFloat64> b = tf.constant( new double[][] {
+   *         {-14.276086, -12.433481, -2.2447076, -1.5775859, 1.8588694}
+   *     });
+   * Operand<TFloat64> result = fops.linalg.matmul(a, b);
+   * // result = {
+   * //     {127.69746,  111.21564,  20.078575,  14.111271,  -16.62731},
+   * //     {-59.547394, -51.861652, -9.362965,  -6.580314,    7.753584},
+   * //     {  3.1950197,  2.7826407, 0.50237054, 0.35306725, -0.4160191}
+   * //  }
+   *
+   * }</pre>
+   *
+   * <p>Note: This is matrix product, not element-wise product.
+   *
+   * @param a an Operand of of type {@code TFloat16}, {@code TFloat32}, {@code TFloat64 }, {@code
+   *     TInt32}. with a {@code rank > 1}
+   * @param b an Operand with same type and rank as {@code a}.
+   * @param <T> the data type of the Operands
+   * @return A Operand of the same type as {@code a} and {@code b} where each inner-most matrix is
+   *     the product of the corresponding matrices in {@code a} and {@code b}. This is the matrix
+   *     product not an element-wise product.
+   * @throws java.lang.IllegalArgumentException If {@code transposeA} and {@code adjointA} , or
+   *     {@code transposeB} and {@code adjointB} are both set to `true`.
+   */
+  public <T extends TNumber> Operand<T> matmul(Operand<T> a, Operand<T> b) {
+    return MatMul.matmul(scope, a, b, false, false, false, false, false, false);
+  }
+
+  /**
+   * Multiplies matrix {@code a} by matrix {@code b}, producing {@code a} * {@code b }.
+   *
+   * <p>The inputs must, following any transpositions, be tensors of {@code rank >= 2} where the
+   * inner 2 dimensions specify valid matrix multiplication dimensions, and any further outer
+   * dimensions specify matching batch size.
+   *
+   * <p>Both matrices must be of the same type. The supported types are: {@code TFloat16}, {@code
+   * TFloat32}, {@code TFloat64}, {@code TInt32}.
+   *
+   * <p>Either matrix can be transposed or adjointed (conjugated and transposed) on the fly by
+   * setting one of the corresponding flag to true. These are false by default.
+   *
+   * <p>Note: This is matrix product, not element-wise product.
+   *
+   * <p>A simple 2-D tensor matrix multiplication:
+   *
+   * <pre>{@code
+   * Operand<TFloat64> a = tf.constant(new double[][] {
+   *         {-8.944851},
+   *         {4.1711287},
+   *         {-0.22380222}
+   *     });
+   * Operand<TFloat64> b = tf.constant( new double[][] {
+   *         {-14.276086, -12.433481, -2.2447076, -1.5775859, 1.8588694}
+   *     });
+   * Operand<TFloat64> result = fops.linalg.matmul(a, b);
+   * // result = {
+   * //     {127.69746,  111.21564,  20.078575,  14.111271,  -16.62731},
+   * //     {-59.547394, -51.861652, -9.362965,  -6.580314,    7.753584},
+   * //     {  3.1950197,  2.7826407, 0.50237054, 0.35306725, -0.4160191}
+   * //  }
+   *
+   * }</pre>
+   *
+   * @param a an Operand of of type {@code TFloat16}, {@code TFloat32}, {@code TFloat64 }, {@code
+   *     TInt32}. with a {@code rank > 1}
+   * @param b an Operand with same type and rank as {@code a}.
+   * @param transposeA If true, {@code a} is transposed before multiplication.
+   * @param transposeB If true, {@code b} is transposed before multiplication
+   * @param <T> the data type of the Operands
+   * @return A Operand of the same type as {@code a} and {@code b} where each inner-most matrix is
+   *     the product of the corresponding matrices in {@code a} and {@code b}. This is the matrix
+   *     product not an element-wise product.
+   * @throws java.lang.IllegalArgumentException If {@code transposeA} and {@code adjointA} , or
+   *     {@code transposeB} and {@code adjointB} are both set to `true`.
+   */
+  public <T extends TNumber> Operand<T> matmul(
+      Operand<T> a, Operand<T> b, boolean transposeA, boolean transposeB) {
+    return MatMul.matmul(scope, a, b, transposeA, transposeB, false, false, false, false);
+  }
+
+  /**
+   * Multiplies matrix {@code a} by matrix {@code b}, producing {@code a} * {@code b }.
+   *
+   * <p>The inputs must, following any transpositions, be tensors of {@code rank >= 2} where the
+   * inner 2 dimensions specify valid matrix multiplication dimensions, and any further outer
+   * dimensions specify matching batch size.
+   *
+   * <p>Both matrices must be of the same type. The supported types are: {@code TFloat16}, {@code
+   * TFloat32}, {@code TFloat64}, {@code TInt32}.
+   *
+   * <p>Either matrix can be transposed or adjointed (conjugated and transposed) on the fly by
+   * setting one of the corresponding flag to true. These are false by default.
+   *
+   * <p>Note: This is matrix product, not element-wise product.
+   *
+   * <p>A simple 2-D tensor matrix multiplication:
+   *
+   * <pre>{@code
+   * Operand<TFloat64> a = tf.constant(new double[][] {
+   *         {-8.944851},
+   *         {4.1711287},
+   *         {-0.22380222}
+   *     });
+   * Operand<TFloat64> b = tf.constant( new double[][] {
+   *         {-14.276086, -12.433481, -2.2447076, -1.5775859, 1.8588694}
+   *     });
+   * Operand<TFloat64> result = fops.linalg.matmul(a, b);
+   * // result = {
+   * //     {127.69746,  111.21564,  20.078575,  14.111271,  -16.62731},
+   * //     {-59.547394, -51.861652, -9.362965,  -6.580314,    7.753584},
+   * //     {  3.1950197,  2.7826407, 0.50237054, 0.35306725, -0.4160191}
+   * //  }
+   *
+   * }</pre>
+   *
+   * @param a an Operand of of type {@code TFloat16}, {@code TFloat32}, {@code TFloat64 }, {@code
+   *     TInt32}. with a {@code rank > 1}
+   * @param b an Operand with same type and rank as {@code a}.
+   * @param transposeA If true, {@code a} is transposed before multiplication.
+   * @param transposeB If True, {@code b} is transposed before multiplication
+   * @param adjointA If true, {@code a} is conjugated and transposed before multiplication.
+   * @param adjointB If true, {@code b} is conjugated and transposed before multiplication.
+   * @param aIsSparse If true, {@code a} is treated as a sparse matrix. Notice, this <em>does not
+   *     support {@link SparseTensor}</em>, it just makes optimizations that assume most values in
+   *     {@code a} are zero.
+   * @param bIsSparse If true, {@code b} is treated as a sparse matrix. Notice, this <em>does not
+   *     support {@link SparseTensor}</em>, it just makes optimizations that assume most values in
+   *     {@code b} are zero.
+   * @param <T> the data type of the Operands
+   * @return A Operand of the same type as {@code a} and {@code b} where each inner-most matrix is
+   *     the product of the corresponding matrices in {@code a} and {@code b}. This is the matrix
+   *     product not an element-wise product.
+   * @throws java.lang.IllegalArgumentException If {@code transposeA} and {@code adjointA} , or
+   *     {@code transposeB} and {@code adjointB} are both set to `true`.
+   */
+  public <T extends TNumber> Operand<T> matmul(
+      Operand<T> a,
+      Operand<T> b,
+      boolean transposeA,
+      boolean transposeB,
+      boolean adjointA,
+      boolean adjointB,
+      boolean aIsSparse,
+      boolean bIsSparse) {
+    return MatMul.matmul(
+        scope, a, b, transposeA, transposeB, adjointA, adjointB, aIsSparse, bIsSparse);
+  }
+}
diff --git a/tensorflow-framework/src/main/java/org/tensorflow/framework/op/MathOps.java b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/MathOps.java
new file mode 100644
index 00000000000..d8b0c7e6775
--- /dev/null
+++ b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/MathOps.java
@@ -0,0 +1,456 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+=======================================================================*/
+package org.tensorflow.framework.op;
+
+import org.tensorflow.Operand;
+import org.tensorflow.framework.op.math.Axes;
+import org.tensorflow.framework.op.math.ConfusionMatrix;
+import org.tensorflow.framework.op.math.L2Normalize;
+import org.tensorflow.framework.op.math.ReduceLogSumExp;
+import org.tensorflow.framework.op.math.TensorDot;
+import org.tensorflow.op.Scope;
+import org.tensorflow.types.TInt32;
+import org.tensorflow.types.TInt64;
+import org.tensorflow.types.family.TFloating;
+import org.tensorflow.types.family.TNumber;
+import org.tensorflow.types.family.TType;
+
+public class MathOps {
+  private final Scope scope;
+
+  private final FrameworkOps frameworkOps;
+
+  /**
+   * Creates Framework math Operations
+   *
+   * @param frameworkOps the TensorFLow framework Ops
+   */
+  MathOps(FrameworkOps frameworkOps) {
+    this.scope = frameworkOps.scope();
+    this.frameworkOps = frameworkOps;
+  }
+
+  /**
+   * Normalizes along dimension axis using an L2 norm.
+   *
+   * @param x the input
+   * @param axis Dimension along which to normalize.
+   * @param <T> the data type for the input and the result
+   * @return the normalized values based on L2 norm
+   */
+  public <T extends TNumber> Operand<T> l2Normalize(Operand<T> x, int[] axis) {
+    return L2Normalize.l2Normalize(scope, x, axis);
+  }
+
+  /**
+   * Computes the confusion matrix from predictions and labels.
+   *
+   * <p>The matrix columns represent the prediction labels and the rows represent the real labels.
+   * The confusion matrix is always a 2-D array of shape {@code [n,n]}, where {@code n} is the
+   * number of valid labels for a given classification task. Both prediction and labels must be 1-D
+   * arrays of the same shape in order for this function to work.
+   *
+   * <p>If {@code numClasses} is null, then {@code numClasses} will be set to one plus the maximum
+   * value in either predictions or labels. Class labels are expected to start at 0. For example, if
+   * {@code numClasses} is 3, then the possible labels would be {@code [0, 1, 2]}.
+   *
+   * <p>If {@code weights} is not null, then each prediction contributes its corresponding weight to
+   * the total value of the confusion matrix cell.
+   *
+   * <p>For example:
+   *
+   * <pre>{@code
+   * fops.math.confusion_matrix(tf.constant(new int[] {1, 2, 4}), tf.constant(new int[] {2, 2, 4})) ==>
+   *     [[0 0 0 0 0]
+   *      [0 0 1 0 0]
+   *      [0 0 1 0 0]
+   *      [0 0 0 0 0]
+   *      [0 0 0 0 1]]
+   * }</pre>
+   *
+   * <p>Note that the possible labels are assumed to be {@code [0, 1, 2, 3, 4]}, resulting in a 5x5
+   * confusion matrix.
+   *
+   * @param labels 1-D Operand of real labels for the classification task.
+   * @param predictions 1-D Operand of predictions for a given classification.
+   * @param <T> Data type of the confusion matrix.
+   * @return An Operand of type {@code type} with shape {@code [n, n]} representing the confusion
+   *     matrix, where {@code n} is the number of possible labels in the classification task.
+   * @throws IllegalArgumentException If both predictions and labels are not 1-D vectors and have
+   *     mismatched shapes, or if {@code weights} is not null and its shape doesn't match {@code
+   *     predictions}.
+   */
+  public <T extends TNumber> Operand<T> confusionMatrix(Operand<T> labels, Operand<T> predictions) {
+    return ConfusionMatrix.confusionMatrix(scope, labels, predictions, null, null);
+  }
+
+  /**
+   * Computes the confusion matrix from predictions and labels.
+   *
+   * <p>The matrix columns represent the prediction labels and the rows represent the real labels.
+   * The confusion matrix is always a 2-D array of shape {@code [n,n]}, where {@code n} is the
+   * number of valid labels for a given classification task. Both prediction and labels must be 1-D
+   * arrays of the same shape in order for this function to work.
+   *
+   * <p>If {@code numClasses} is null, then {@code numClasses} will be set to one plus the maximum
+   * value in either predictions or labels. Class labels are expected to start at 0. For example, if
+   * {@code numClasses} is 3, then the possible labels would be {@code [0, 1, 2]}.
+   *
+   * <p>If {@code weights} is not null, then each prediction contributes its corresponding weight to
+   * the total value of the confusion matrix cell.
+   *
+   * <p>For example:
+   *
+   * <pre>{@code
+   * fops.math.confusion_matrix(tf.constant(new int[] {1, 2, 4}), tf.constant(new int[] {2, 2, 4})) ==>
+   *     [[0 0 0 0 0]
+   *      [0 0 1 0 0]
+   *      [0 0 1 0 0]
+   *      [0 0 0 0 0]
+   *      [0 0 0 0 1]]
+   * }</pre>
+   *
+   * <p>Note that the possible labels are assumed to be {@code [0, 1, 2, 3, 4]}, resulting in a 5x5
+   * confusion matrix.
+   *
+   * @param labels 1-D Operand of real labels for the classification task.
+   * @param predictions 1-D Operand of predictions for a given classification.
+   * @param weights An optional Operand whose shape matches {@code predictions}.
+   * @param <T> Data type of the confusion matrix.
+   * @return An Operand of type {@code type} with shape {@code [n, n]} representing the confusion
+   *     matrix, where {@code n} is the number of possible labels in the classification task.
+   * @throws IllegalArgumentException If both predictions and labels are not 1-D vectors and have
+   *     mismatched shapes, or if {@code weights} is not null and its shape doesn't match {@code
+   *     predictions}.
+   */
+  public <T extends TNumber> Operand<T> confusionMatrix(
+      Operand<T> labels, Operand<T> predictions, Operand<T> weights) {
+    return ConfusionMatrix.confusionMatrix(scope, labels, predictions, weights, null);
+  }
+
+  /**
+   * Computes the confusion matrix from predictions and labels.
+   *
+   * <p>The matrix columns represent the prediction labels and the rows represent the real labels.
+   * The confusion matrix is always a 2-D array of shape {@code [n,n]}, where {@code n} is the
+   * number of valid labels for a given classification task. Both prediction and labels must be 1-D
+   * arrays of the same shape in order for this function to work.
+   *
+   * <p>If {@code numClasses} is null, then {@code numClasses} will be set to one plus the maximum
+   * value in either predictions or labels. Class labels are expected to start at 0. For example, if
+   * {@code numClasses} is 3, then the possible labels would be {@code [0, 1, 2]}.
+   *
+   * <p>If {@code weights} is not null, then each prediction contributes its corresponding weight to
+   * the total value of the confusion matrix cell.
+   *
+   * <p>For example:
+   *
+   * <pre>{@code
+   * fops.math.confusion_matrix(tf.constant(new int[] {1, 2, 4}), tf.constant(new int[] {2, 2, 4})) ==>
+   *     [[0 0 0 0 0]
+   *      [0 0 1 0 0]
+   *      [0 0 1 0 0]
+   *      [0 0 0 0 0]
+   *      [0 0 0 0 1]]
+   * }</pre>
+   *
+   * <p>Note that the possible labels are assumed to be {@code [0, 1, 2, 3, 4]}, resulting in a 5x5
+   * confusion matrix.
+   *
+   * @param labels 1-D Operand of real labels for the classification task.
+   * @param predictions 1-D Operand of predictions for a given classification.
+   * @param weights An optional Operand whose shape matches {@code predictions}.
+   * @param numClasses The possible number of labels the classification task can have. If this value
+   *     is null, it will be calculated using both predictions and labels.
+   * @param <T> Data type of the confusion matrix.
+   * @return An Operand of type {@code type} with shape {@code [n, n]} representing the confusion
+   *     matrix, where {@code n} is the number of possible labels in the classification task.
+   * @throws IllegalArgumentException If both predictions and labels are not 1-D vectors and have
+   *     mismatched shapes, or if {@code weights} is not null and its shape doesn't match {@code
+   *     predictions}.
+   */
+  public <T extends TNumber> Operand<T> confusionMatrix(
+      Operand<T> labels, Operand<T> predictions, Operand<T> weights, Operand<TInt64> numClasses) {
+    return ConfusionMatrix.confusionMatrix(scope, labels, predictions, weights, numClasses);
+  }
+
+  /**
+   * Creates an Operand that has all axes contained in the Operand's shape.
+   *
+   * @param op the Operand
+   * @return an Operand that has all axes contained in the Operand's shape..
+   */
+  public Operand<TInt32> allAxes(Operand<? extends TType> op) {
+    return Axes.allAxes(scope, op);
+  }
+
+  /**
+   * Tensor contraction of a and b along specified axes and outer product.
+   *
+   * <p>Tensordot (also known as tensor contraction) sums the product of elements from {@code a} and
+   * {@code b} over the indices specified by {@code a_axes} and {@code b_axes}. The lists {@code
+   * a_axes} and {@code b_axes} specify those pairs of axes along which to contract the tensors. The
+   * axis {@code a_axes[i]} of {@code a} must have the same dimension as axis {@code b_axes[i]} of
+   * {@code b} for all {@code i} in {@code range(0, len(a_axes))}. The lists {@code a_axes} and
+   * {@code b_axes} must have identical length and consist of unique integers that specify valid
+   * axes for each of the tensors. Additionally outer product is supported by passing {@code
+   * axes=0}.
+   *
+   * <p>This operation corresponds to {@code numpy.tensordot(a, b, axes)}.
+   *
+   * <p>Example 1: When {@code a} and {@code b} are matrices (order 2), the case {@code axes = 1} is
+   * equivalent to matrix multiplication.
+   *
+   * <p>Example 2: When {@code a} and {@code b} are matrices (order 2), the case {@code axes = [[1],
+   * [0]]} is equivalent to matrix multiplication.
+   *
+   * <p>Example 3: When {@code a} and {@code b} are matrices (order 2), the case {@code axes=0}
+   * gives the outer product, a tensor of order 4.
+   *
+   * <p>Example 4: Suppose that <i>a<sub>ijk</sub></i> and <i>b<sub>lmn</sub></i> represent two
+   * tensors of order 3. Then, {@code contract(a, b, [[0], [2]])} is the order 4 tensor
+   * <i>c<sub>jklm</sub></i> whose entry corresponding to the indices <i>(j,k,l,m)</i> is given by:
+   * <i> c<sub>jklm</sub> = &Sigma;<sub>i</sub> a<sub>ijk</sub> b<sub>lmi</sub> </i>.
+   *
+   * <p>In general, {@code order(c) = order(a) + order(b) - 2*len(axes[0])}.
+   *
+   * @param a {@code Operand} of type {@code TFloat32} or {@code TFloat64}.
+   * @param b {@code Operand} with the same type as {@code a}.
+   * @param axis sum over the last N axes of a and the first N axes of b in order. If {@code
+   *     axis=0}, computes the outer product between {@code a} and {@code b}.
+   * @param <T> the datatype of the Operands, must be either TFloat32 or TFloat64
+   * @return A {@code Operand} with the same type as {@code a}.
+   * @throws IllegalArgumentException if a is not a float32 or float64 data type and if a and b are
+   *     not the same data type
+   */
+  public <T extends TFloating> Operand<T> tensordot(Operand<T> a, Operand<T> b, int axis) {
+    return TensorDot.tensordot(scope, a, b, axis);
+  }
+
+  /**
+   * Tensor contraction of a and b along specified axes and outer product.
+   *
+   * <p>Tensordot (also known as tensor contraction) sums the product of elements from {@code a} and
+   * {@code b} over the indices specified by {@code a_axes} and {@code b_axes}. The lists {@code
+   * a_axes} and {@code b_axes} specify those pairs of axes along which to contract the tensors. The
+   * axis {@code a_axes[i]} of {@code a} must have the same dimension as axis {@code b_axes[i]} of
+   * {@code b} for all {@code i} in {@code range(0, len(a_axes))}. The lists {@code a_axes} and
+   * {@code b_axes} must have identical length and consist of unique integers that specify valid
+   * axes for each of the tensors. Additionally outer product is supported by passing {@code
+   * axes=0}.
+   *
+   * <p>This operation corresponds to {@code numpy.tensordot(a, b, axes)}.
+   *
+   * <p>Example 1: When {@code a} and {@code b} are matrices (order 2), the case {@code axes = 1} is
+   * equivalent to matrix multiplication.
+   *
+   * <p>Example 2: When {@code a} and {@code b} are matrices (order 2), the case {@code axes = [[1],
+   * [0]]} is equivalent to matrix multiplication.
+   *
+   * <p>Example 3: When {@code a} and {@code b} are matrices (order 2), the case {@code axes=0}
+   * gives the outer product, a tensor of order 4.
+   *
+   * <p>Example 4: Suppose that <i>a<sub>ijk</sub></i> and <i>b<sub>lmn</sub></i> represent two
+   * tensors of order 3. Then, {@code contract(a, b, [[0], [2]])} is the order 4 tensor
+   * <i>c<sub>jklm</sub></i> whose entry corresponding to the indices <i>(j,k,l,m)</i> is given by:
+   *
+   * <p><i> c<sub>jklm</sub> = &Sigma;<sub>i</sub> a<sub>ijk</sub> b<sub>lmi</sub> </i>.
+   *
+   * <p>In general, {@code order(c) = order(a) + order(b) - 2*len(axes[0])}.
+   *
+   * <p>
+   *
+   * @param a {@code Operand} of type {@code TFloat32} or {@code TFloat64}.
+   * @param b {@code Operand} with the same type as {@code a}.
+   * @param axes If axes is a scalar, sum over the last N axes of a and the first N axes of b in
+   *     order. If axes is a list, the first and second row contain the set of unique integers
+   *     specifying axes along which the contraction is computed, for {@code a} and {@code b},
+   *     respectively. The number of axes for {@code a} and {@code b} must be equal. If {@code
+   *     axis=0}, computes the outer product between {@code a} and {@code b}.
+   * @param <T> the datatype of the Operands, must be either TFloat32 or TFloat64
+   * @return A {@code Operand} with the same type as {@code a}.
+   * @throws IllegalArgumentException if a is not a float32 or float64 data type and if a and b are
+   *     not the same data type
+   */
+  public <T extends TFloating> Operand<T> tensordot(
+      Operand<T> a, Operand<T> b, Operand<TInt32> axes) {
+    return TensorDot.tensordot(scope, a, b, axes);
+  }
+
+  /**
+   * Tensor contraction of a and b along specified axes and outer product.
+   *
+   * <p>Tensordot (also known as tensor contraction) sums the product of elements from {@code a} and
+   * {@code b} over the indices specified by {@code a_axes} and {@code b_axes}. The lists {@code
+   * a_axes} and {@code b_axes} specify those pairs of axes along which to contract the tensors. The
+   * axis {@code a_axes[i]} of {@code a} must have the same dimension as axis {@code b_axes[i]} of
+   * {@code b} for all {@code i} in {@code range(0, len(a_axes))}. The lists {@code a_axes} and
+   * {@code b_axes} must have identical length and consist of unique integers that specify valid
+   * axes for each of the tensors. Additionally outer product is supported by passing {@code
+   * axes=0}.
+   *
+   * <p>This operation corresponds to {@code numpy.tensordot(a, b, axes)}.
+   *
+   * <p>Example 1: When {@code a} and {@code b} are matrices (order 2), the case {@code axes = 1} is
+   * equivalent to matrix multiplication.
+   *
+   * <p>Example 2: When {@code a} and{@code b} are matrices (order 2), the case {@code axes = [[1],
+   * [0]]} is equivalent to matrix multiplication.
+   *
+   * <p>Example 3: When {@code a} and {@code b} are matrices (order 2), the case {@code axes=0}
+   * gives the outer product, a tensor of order 4.
+   *
+   * <p>Example 4: Suppose that <i>a<sub>ijk</sub></i> and <i>b<sub>lmn</sub></i> represent two
+   * tensors of order 3. Then, {@code contract(a, b, [[0], [2]])} is the order 4 tensor
+   * <i>c<sub>jklm</sub></i> whose entry corresponding to the indices <i>(j,k,l,m)</i> is given by:
+   *
+   * <p><i> c<sub>jklm</sub> = &Sigma;<sub>i</sub> a<sub>ijk</sub> b<sub>lmi</sub> </i>.
+   *
+   * <p>In general, {@code order(c) = order(a) + order(b) - 2*len(axes[0])}.
+   *
+   * <p>
+   *
+   * @param a {@code Operand} of type {@code TFloat32} or {@code TFloat64}.
+   * @param b {@code Operand} with the same type as {@code a}.
+   * @param axes the first and second row contain the set of unique integers specifying axes along
+   *     which the contraction is computed, for {@code a} and {@code b}, respectively. The number of
+   *     axes for {@code a} and {@code b} must be equal. I
+   * @param <T> the datatype of the Operands, must be either TFloat32 or TFloat64
+   * @return A {@code Operand} with the same type as {@code a}.
+   * @throws IllegalArgumentException if a is not a float32 or float64 data type and if a and b are
+   *     not the same data type
+   */
+  public <T extends TFloating> Operand<T> tensordot(Operand<T> a, Operand<T> b, int[] axes) {
+    return TensorDot.tensordot(scope, a, b, axes);
+  }
+
+  /**
+   * Tensor contraction of a and b along specified axes and outer product.
+   *
+   * <p>Tensordot (also known as tensor contraction) sums the product of elements from {@code a} and
+   * {@code b} over the indices specified by {@code a_axes} and {@code b_axes}. The lists {@code
+   * a_axes} and {@code b_axes} specify those pairs of axes along which to contract the tensors. The
+   * axis {@code a_axes[i]} of {@code a} must have the same dimension as axis {@code b_axes[i]} of
+   * {@code b} for all {@code i} in {@code range(0, len(a_axes))}. The lists {@code a_axes} and
+   * {@code b_axes} must have identical length and consist of unique integers that specify valid
+   * axes for each of the tensors. Additionally outer product is supported by passing {@code
+   * axes=0}.
+   *
+   * <p>This operation corresponds to {@code numpy.tensordot(a, b, axes)}.
+   *
+   * <p>Example 1: When {@code a} and {@code b} are matrices (order 2), the case {@code axes = 1} is
+   * equivalent to matrix multiplication.
+   *
+   * <p>Example 2: When {@code a} and{@code b} are matrices (order 2), the case {@code axes = [[1],
+   * [0]]} is equivalent to matrix multiplication.
+   *
+   * <p>Example 3: When {@code a} and {@code b} are matrices (order 2), the case {@code axes=0}
+   * gives the outer product, a tensor of order 4.
+   *
+   * <p>Example 4: Suppose that <i>a<sub>ijk</sub></i> and <i>b<sub>lmn</sub></i> represent two
+   * tensors of order 3. Then, {@code contract(a, b, [[0], [2]])} is the order 4 tensor
+   * <i>c<sub>jklm</sub></i> whose entry corresponding to the indices <i>(j,k,l,m)</i> is given by:
+   *
+   * <p><i> c<sub>jklm</sub> = &Sigma;<sub>i</sub> a<sub>ijk</sub> b<sub>lmi</sub> </i>.
+   *
+   * <p>In general, {@code order(c) = order(a) + order(b) - 2*len(axes[0])}.
+   *
+   * <p>
+   *
+   * @param a {@code Operand} of type {@code TFloat32} or {@code TFloat64}.
+   * @param b {@code Operand} with the same type as {@code a}.
+   * @param axes the first and second row contain the set of unique integers specifying axes along
+   *     which the contraction is computed, for {@code a} and {@code b}, respectively. The number of
+   *     axes for {@code a} and {@code b} must be equal. I
+   * @param <T> the datatype of the Operands, must be either TFloat32 or TFloat64
+   * @return A {@code Operand} with the same type as {@code a}.
+   * @throws IllegalArgumentException if a is not a float32 or float64 data type and if a and b are
+   *     not the same data type
+   */
+  public <T extends TFloating> Operand<T> tensordot(Operand<T> a, Operand<T> b, int[][] axes) {
+    return TensorDot.tensordot(scope, a, b, axes);
+  }
+
+  /**
+   * Tensor contraction of a and b along specified axes and outer product.
+   *
+   * <p>Tensordot (also known as tensor contraction) sums the product of elements from {@code a} and
+   * {@code b} over the indices specified by {@code a_axes} and {@code b_axes}. The lists {@code
+   * a_axes} and {@code b_axes} specify those pairs of axes along which to contract the tensors. The
+   * axis {@code a_axes[i]} of {@code a} must have the same dimension as axis {@code b_axes[i]} of
+   * {@code b} for all {@code i} in {@code range(0, len(a_axes))}. The lists {@code a_axes} and
+   * {@code b_axes} must have identical length and consist of unique integers that specify valid
+   * axes for each of the tensors. Additionally outer product is supported by passing {@code
+   * axes=0}.
+   *
+   * <p>This operation corresponds to {@code numpy.tensordot(a, b, axes)}.
+   *
+   * <p>Example 1: When {@code a} and {@code b} are matrices (order 2), the case {@code axes = 1} is
+   * equivalent to matrix multiplication.
+   *
+   * <p>Example 2: When {@code a} and{@code b} are matrices (order 2), the case {@code axes = [[1],
+   * [0]]} is equivalent to matrix multiplication.
+   *
+   * <p>Example 3: When {@code a} and {@code b} are matrices (order 2), the case {@code axes=0}
+   * gives the outer product, a tensor of order 4.
+   *
+   * <p>Example 4: Suppose that <i>a<sub>ijk</sub></i> and <i>b<sub>lmn</sub></i> represent two
+   * tensors of order 3. Then, {@code contract(a, b, [[0], [2]])} is the order 4 tensor
+   * <i>c<sub>jklm</sub></i> whose entry corresponding to the indices <i>(j,k,l,m)</i> is given by:
+   *
+   * <p><i> c<sub>jklm</sub> = &Sigma;<sub>i</sub> a<sub>ijk</sub> b<sub>lmi</sub> </i>.
+   *
+   * <p>In general, {@code order(c) = order(a) + order(b) - 2*len(axes[0])}.
+   *
+   * <p>
+   *
+   * @param a {@code Operand} of type {@code TFloat32} or {@code TFloat64}.
+   * @param b {@code Operand} with the same type as {@code a}.
+   * @param aAxis axes for the a Operand
+   * @param bAxis axes for the b Operand
+   * @param <T> the datatype of the Operands, must be either TFloat32 or TFloat64
+   * @return A {@code Operand} with the same type as {@code a}.
+   * @throws IllegalArgumentException if a is not a float32 or float64 data type and if a and b are
+   *     not the same data type
+   */
+  public <T extends TFloating> Operand<T> tensordot(
+      Operand<T> a, Operand<T> b, Operand<TInt32> aAxis, Operand<TInt32> bAxis) {
+    return TensorDot.tensordot(scope, a, b, aAxis, bAxis);
+  }
+
+  /**
+   * Computes log(sum(exp(elements across dimensions of a tensor))). Reduces {@code input_tensor}
+   * along the dimensions given in {@code axes}.
+   *
+   * <p>Reduces {@code input} along the dimensions given in {@code axes}. Unless {@code keepdims} is
+   * true, the rank of the tensor is reduced by 1 for each of the entries in {@code axes}, which
+   * must be unique. If {@code keepdims} is true, the reduced dimensions are retained with length 1.
+   * If {@code axes} has no entries, all dimensions are reduced, and a tensor with a single element
+   * is returned. This function is more numerically stable than {@code log(sum(exp(input)))}. It
+   * avoids overflows caused by taking the exp of large inputs and underflows caused by taking the
+   * log of small inputs.
+   *
+   * @param input The tensor to reduce.
+   * @param axes The dimensions to reduce. If null, reduces all dimensions. Must be in the range
+   *     {@code [-rank(input_tensor), rank(input_tensor)]}.
+   * @param keepDims If true, retains reduced dimensions with length 1.
+   * @param <T> the data type for the input and the result
+   * @return The reduced tensor.
+   */
+  public <T extends TFloating> Operand<T> reduceLogSumExp(
+      Operand<T> input, int[] axes, boolean keepDims) {
+    return ReduceLogSumExp.reduceLogSumExp(scope, input, axes, keepDims);
+  }
+}
diff --git a/tensorflow-framework/src/main/java/org/tensorflow/framework/op/NnOps.java b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/NnOps.java
new file mode 100644
index 00000000000..96f023ffedf
--- /dev/null
+++ b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/NnOps.java
@@ -0,0 +1,193 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+=======================================================================*/
+package org.tensorflow.framework.op;
+
+import org.tensorflow.Operand;
+import org.tensorflow.framework.op.nn.SigmoidCrossEntropyWithLogits;
+import org.tensorflow.framework.op.nn.SoftmaxCrossEntropyWithLogits;
+import org.tensorflow.framework.op.nn.SparseSoftmaxCrossEntropyWithLogits;
+import org.tensorflow.op.Scope;
+import org.tensorflow.types.family.TNumber;
+
+/**
+ * Creates Framework nerual network Operations
+ *
+ * <p>These are higher level ops that may invoke core ops. Higher level Ops may perform the
+ * operation solely in the TensorFlow framework or do preprocessing of the Operands before invoking
+ * a core level Op.
+ *
+ * <p>{@link FrameworkOps}
+ */
+public class NnOps {
+  private final Scope scope;
+
+  private final FrameworkOps frameworkOps;
+
+  /**
+   * Creates Framework {@code nn} Operations
+   *
+   * @param frameworkOps the TensorFLow framework Ops
+   */
+  NnOps(FrameworkOps frameworkOps) {
+    this.scope = frameworkOps.scope();
+    this.frameworkOps = frameworkOps;
+  }
+
+  /**
+   * Computes sigmoid cross entropy given {@code logits}.
+   *
+   * <p>Measures the probability error in discrete classification tasks in which each class is
+   * independent and not mutually exclusive. For instance, one could perform multilabel
+   * classification where a picture can contain both an elephant and a dog at the same time.
+   *
+   * <p>For brevity, let {@code x = logits}, {@code z = labels}. The logistic loss in pseudo-code is
+   *
+   * <pre>
+   *  z * -log(sigmoid(x)) + (1 - z) * -log(1 - sigmoid(x))
+   *   = z * -log(1 / (1 + exp(-x))) + (1 - z) * -log(exp(-x) / (1 + exp(-x)))
+   *   = z * log(1 + exp(-x)) + (1 - z) * (-log(exp(-x)) + log(1 + exp(-x)))
+   *   = z * log(1 + exp(-x)) + (1 - z) * (x + log(1 + exp(-x))
+   *   = (1 - z) * x + log(1 + exp(-x))
+   *   = x - x * z + log(1 + exp(-x))
+   *  </pre>
+   *
+   * <p>For {@code x < 0}, to avoid overflow in {@code exp(-x)}, we reformulate the above
+   *
+   * <pre>
+   *  x - x * z + log(1 + exp(-x))
+   *   = log(exp(x)) - x * z + log(1 + exp(-x))
+   *   = - x * z + log(1 + exp(x))
+   *  </pre>
+   *
+   * <p>Hence, to ensure stability and avoid overflow, the implementation uses this equivalent
+   * formulation
+   *
+   * <pre>
+   *    max(x, 0) - x * z + log(1 + exp(-abs(x)))
+   *  </pre>
+   *
+   * <p>{@code logits} and {@code labels} must have the same type and shape.
+   *
+   * <p>
+   *
+   * @param labels the labels
+   * @param logits the logits of type float32 or float64
+   * @param <T> the type of labels and logits
+   * @return the component-wise logistic losses.
+   * @throws IllegalArgumentException if logits and labels do not have the same shape
+   */
+  public <T extends TNumber> Operand<T> sigmoidCrossEntropyWithLogits(
+      Operand<T> labels, Operand<T> logits) {
+    return SigmoidCrossEntropyWithLogits.sigmoidCrossEntropyWithLogits(scope, labels, logits);
+  }
+
+  /**
+   * Computes softmax cross entropy between {@code logits} and {@code labels}.
+   *
+   * <p>Measures the probability error in discrete classification tasks in which the classes are
+   * mutually exclusive (each entry is in exactly one class). For example, each CIFAR-10 image is
+   * labeled with one and only one label: an image can be a dog or a truck, but not both.
+   *
+   * <p><b>NOTE:</b>
+   *
+   * <p>While the classes are mutually exclusive, their probabilities need not be. All that is
+   * required is that each row of {@code labels} is a valid probability distribution. If they are
+   * not, the computation of the gradient will be incorrect.
+   *
+   * <p>If using exclusive {@code labels} (wherein one and only one class is true at a time), see
+   * {@link org.tensorflow.op.NnOps#sparseSoftmaxCrossEntropyWithLogits}
+   *
+   * <p>Usage:
+   *
+   * <pre>
+   *    Operand&lt;TFloat32&gt; logits =
+   *        tf.constant(new float[][] {{4.0F, 2.0F, 1.0F}, {0.0F, 5.0F, 1.0F}} );
+   *    Operand&lt;TFloat32&gt; labels =
+   *        tf.constant(new float[][] {{1.0F, 0.0F, 0.0F}, {0.0F, 0.8F, 0.2F}} );
+   *    Operand&lt;TFloat32&gt; output =
+   *        tf.nn.softmaxCrossEntropyWithLogits(labels, logits, -1);
+   *    // output Shape = [2]
+   *    // dataType = FLOAT (1)
+   *    // values { 0.169846, 0.824745 }
+   *  </pre>
+   *
+   * <p>Backpropagation will happen into both {@code logits} and {@code labels}. To disallow
+   * backpropagation into {@code labels}, pass label tensors through {@code tf.stopGradient} before
+   * feeding it to this function.
+   *
+   * @param labels Each vector along the class dimension should hold a valid probability
+   *     distribution e.g. for the case in which labels are of shape {@code [batch_size,
+   *     num_classes] }, each row of {@code labels[i]} must be a valid probability distribution.
+   * @param logits Per-label activations, typically a linear output. These activation energies are
+   *     interpreted as unnormalized log probabilities.
+   * @param axis The class dimension. -1 is the last dimension.
+   * @param <T> the number type of the operands
+   * @param <U> the data type for the labels.
+   * @return the softmax cross entropy loss. Its type is the same as {@code logits} and its shape is
+   *     the same as {@code labels} except that it does not have the last dimension of {@code
+   *     labels}.
+   */
+  public <T extends TNumber, U extends TNumber> Operand<T> softmaxCrossEntropyWithLogits(
+      Operand<U> labels, Operand<T> logits, int axis) {
+    return SoftmaxCrossEntropyWithLogits.softmaxCrossEntropyWithLogits(scope, labels, logits, axis);
+  }
+
+  /**
+   * Computes sparse softmax cross entropy between {@code logits} and {@code labels}.
+   *
+   * <p>Measures the probability error in discrete classification tasks in which the classes are
+   * mutually exclusive (each entry is in exactly one class). For example, each CIFAR-10 image is
+   * labeled with one and only one label: an image can be a dog or a truck, but not both.
+   *
+   * <p><b>NOTE:</b>
+   *
+   * <p>For this operation, the probability of a given label is considered exclusive. That is, soft
+   * classes are not allowed, and the {@code labels} vector must provide a single specific index for
+   * the true class for each row of {@code logits} (each minibatch entry). For soft softmax
+   * classification with a probability distribution for each entry, {@link
+   * org.tensorflow.op.NnOps#softmaxCrossEntropyWithLogits}.
+   *
+   * <p><b>WARNING:</b>
+   *
+   * <p>This op expects unscaled logits, since it performs a {@code softmax} on {@code logits }
+   * internally for efficiency. Do not call this op with the output of {@code softmax}, as it will
+   * produce incorrect results.
+   *
+   * <p>A common use case is to have logits of shape {@code [batchSize, numClasses]} and have labels
+   * of shape {@code [batchSize]}, but higher dimensions are supported, in which case the {@code
+   * dim}-th dimension is assumed to be of size {@code numClasses}. {@code logits} must have the
+   * {@code dataType} of {@code TFloat16}, {@code TFloat32} , or {@code TFloat64}, and {@code
+   * labels} must have the dtype of {@code TInt32} or {@code TInt64}.
+   *
+   * @param labels {@code Tensor} of shape {@code [d_0, d_1, ..., d_{r-1}]} (where {@code r } is
+   *     rank of {@code labels} and result) and the dataType is {@code TInt32} or {@code TInt64}.
+   *     Each entry in {@code labels} must be an index in {@code [0, numClasses)}. Other values will
+   *     raise an exception when this op is run on CPU, and return {@code NaN} for corresponding
+   *     loss and gradient rows on GPU.
+   * @param logits Per-label activations (typically a linear output) of shape {@code [d_0, d_1, ...,
+   *     d_{r-1}, numClasses]} and dataType of {@code TFloat16}, {@code TFloat32}, or {@code
+   *     TFloat64}. These activation energies are interpreted as unnormalized log probabilities.
+   * @param <U> the data type for the labels
+   * @param <T> the data tyoe for the loss and logits.
+   * @return the loss
+   * @throws IllegalArgumentException If logits are scalars (need to have {@code rank >= 1}) or if
+   *     the rank of the labels is not equal to the rank of the logits minus one.
+   */
+  public <T extends TNumber, U extends TNumber> Operand<T> sparseSoftmaxCrossEntropyWithLogits(
+      Operand<U> labels, Operand<T> logits) {
+    return SparseSoftmaxCrossEntropyWithLogits.sparseSoftmaxCrossEntropyWithLogits(
+        scope, labels, logits);
+  }
+}
diff --git a/tensorflow-framework/src/main/java/org/tensorflow/framework/op/SetOps.java b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/SetOps.java
new file mode 100644
index 00000000000..832f0a2892c
--- /dev/null
+++ b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/SetOps.java
@@ -0,0 +1,101 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+=======================================================================*/
+package org.tensorflow.framework.op;
+
+import org.tensorflow.Operand;
+import org.tensorflow.framework.op.sets.Sets;
+import org.tensorflow.op.Scope;
+import org.tensorflow.types.family.TNumber;
+
+/** Creates Framework set Operations */
+public class SetOps {
+
+  private final Scope scope;
+
+  private final FrameworkOps frameworkOps;
+
+  /**
+   * Creates Framework {@code nn} Operations
+   *
+   * @param frameworkOps the TensorFLow framework Ops
+   */
+  SetOps(FrameworkOps frameworkOps) {
+    this.scope = frameworkOps.scope();
+    this.frameworkOps = frameworkOps;
+  }
+
+  /**
+   * Computes set difference of elements in last dimension of <code>a</code> and <code>b</code> with
+   * <code>aMinusB</code> set to true.
+   *
+   * <p>All but the last dimension of <code>a</code> and <code>b</code> must match
+   *
+   * @param a The first operand representing set <code>a</code>
+   * @param b The other operand representing set <code>b</code>
+   * @param <T>the data type for the sets
+   * @return An Operand with the same rank as <code>a</code> and <code>b</code>, and all but the
+   *     last dimension the * same. Elements along the last dimension contain the results of the set
+   *     operation.
+   */
+  public <T extends TNumber> Operand<T> difference(Operand<T> a, Operand<T> b) {
+
+    return Sets.difference(scope, a, b, true);
+  }
+
+  /**
+   * Computes set difference of elements in last dimension of <code>a</code> and <code>b</code>.
+   *
+   * <p>All but the last dimension of <code>a</code> and <code>b</code> must match
+   *
+   * @param a The first operand representing set <code>a</code>
+   * @param b The other operand representing set <code>b</code>
+   * @param aMinusB whether to subtract b from a, vs vice versa.
+   * @param <T>the data type for the sets
+   * @return An Operand with the same rank as <code>a</code> and <code>b</code>, and all but the
+   *     last dimension the * same. Elements along the last dimension contain the results of the set
+   *     operation.
+   */
+  public <T extends TNumber> Operand<T> difference(Operand<T> a, Operand<T> b, boolean aMinusB) {
+    return Sets.difference(scope, a, b, aMinusB);
+  }
+
+  /**
+   * Computes set union of elements in last dimension of <code>a</code> and <code>b</code>.
+   *
+   * @param a The first operand representing set <code>a</code>
+   * @param b The other operand representing set <code>b</code>
+   * @param <T>the data type for the sets
+   * @return An Operand with the same rank as <code>a</code> and <code>b</code>, and all but the
+   *     last dimension the * same. Elements along the last dimension contain the results of the set
+   *     operation.
+   */
+  public <T extends TNumber> Operand<T> union(Operand<T> a, Operand<T> b) {
+    return Sets.union(scope, a, b);
+  }
+
+  /**
+   * Computes set intersection of elements in last dimension of <code>a</code> and <code>b</code>.
+   *
+   * @param a The first operand representing set <code>a</code>
+   * @param b The other operand representing set <code>b</code>
+   * @param <T>the data type for the sets
+   * @return An Operand with the same rank as <code>a</code> and <code>b</code>, and all but the
+   *     last dimension the * same. Elements along the last dimension contain the results of the set
+   *     operation.
+   */
+  public <T extends TNumber> Operand<T> intersection(Operand<T> a, Operand<T> b) {
+    return Sets.intersection(scope, a, b);
+  }
+}
diff --git a/tensorflow-framework/src/main/java/org/tensorflow/framework/op/linalg/MatMul.java b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/linalg/MatMul.java
new file mode 100644
index 00000000000..a156cdfd2b4
--- /dev/null
+++ b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/linalg/MatMul.java
@@ -0,0 +1,289 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+=======================================================================*/
+package org.tensorflow.framework.op.linalg;
+
+import org.tensorflow.Operand;
+import org.tensorflow.framework.utils.SparseTensor;
+import org.tensorflow.ndarray.Shape;
+import org.tensorflow.op.Scope;
+import org.tensorflow.op.dtypes.Cast;
+import org.tensorflow.op.math.Conj;
+import org.tensorflow.op.sparse.SparseMatMul;
+import org.tensorflow.op.train.BatchMatMul;
+import org.tensorflow.types.TBfloat16;
+import org.tensorflow.types.TFloat32;
+import org.tensorflow.types.TInt32;
+import org.tensorflow.types.family.TFloating;
+import org.tensorflow.types.family.TNumber;
+
+/** Multiplication matrix operations */
+public class MatMul {
+
+  /**
+   * Multiplies matrix {@code a} by matrix {@code b}, producing {@code a} * {@code b }.
+   *
+   * <p>The inputs must, following any transpositions, be tensors of {@code rank >= 2} where the
+   * inner 2 dimensions specify valid matrix multiplication dimensions, and any further outer
+   * dimensions specify matching batch size.
+   *
+   * <p>Both matrices must be of the same type. The supported types are: {@code TFloat16}, {@code
+   * TFloat32}, {@code TFloat64}, {@code TInt32}.
+   *
+   * <p>Either matrix can be transposed or adjointed (conjugated and transposed) on the fly by
+   * setting one of the corresponding flag to true. These are false by default.
+   *
+   * <p>A simple 2-D tensor matrix multiplication:
+   *
+   * <pre>{@code
+   * Operand<TFloat64> a = tf.constant(new double[][] {
+   *         {-8.944851},
+   *         {4.1711287},
+   *         {-0.22380222}
+   *     });
+   * Operand<TFloat64> b = tf.constant( new double[][] {
+   *         {-14.276086, -12.433481, -2.2447076, -1.5775859, 1.8588694}
+   *     });
+   * Operand<TFloat64> result = fops.linalg.matmul(a, b);
+   * // result = {
+   * //     {127.69746,  111.21564,  20.078575,  14.111271,  -16.62731},
+   * //     {-59.547394, -51.861652, -9.362965,  -6.580314,    7.753584},
+   * //     {  3.1950197,  2.7826407, 0.50237054, 0.35306725, -0.4160191}
+   * //  }
+   *
+   * }</pre>
+   *
+   * <p>Note: This is matrix product, not element-wise product.
+   *
+   * @param scope The TensorFlow scope
+   * @param a an Operand of of type {@code TFloat16}, {@code TFloat32}, {@code TFloat64 }, {@code
+   *     TInt32}. with a {@code rank > 1}
+   * @param b an Operand with same type and rank as {@code a}.
+   * @param <T> the data type of the Operands
+   * @return A Operand of the same type as {@code a} and {@code b} where each inner-most matrix is
+   *     the product of the corresponding matrices in {@code a} and {@code b}. This is the matrix
+   *     product not an element-wise product.
+   * @throws java.lang.IllegalArgumentException If {@code transposeA} and {@code adjointA} , or
+   *     {@code transposeB} and {@code adjointB} are both set to `true`.
+   */
+  public static <T extends TNumber> Operand<T> matmul(Scope scope, Operand<T> a, Operand<T> b) {
+    return matmul(scope, a, b, false, false, false, false, false, false);
+  }
+
+  /**
+   * Multiplies matrix {@code a} by matrix {@code b}, producing {@code a} * {@code b }.
+   *
+   * <p>The inputs must, following any transpositions, be tensors of {@code rank >= 2} where the
+   * inner 2 dimensions specify valid matrix multiplication dimensions, and any further outer
+   * dimensions specify matching batch size.
+   *
+   * <p>Both matrices must be of the same type. The supported types are: {@code TFloat16}, {@code
+   * TFloat32}, {@code TFloat64}, {@code TInt32}.
+   *
+   * <p>Either matrix can be transposed or adjointed (conjugated and transposed) on the fly by
+   * setting one of the corresponding flag to true. These are false by default.
+   *
+   * <p>Note: This is matrix product, not element-wise product.
+   *
+   * <p>A simple 2-D tensor matrix multiplication:
+   *
+   * <pre>{@code
+   * Operand<TFloat64> a = tf.constant(new double[][] {
+   *         {-8.944851},
+   *         {4.1711287},
+   *         {-0.22380222}
+   *     });
+   * Operand<TFloat64> b = tf.constant( new double[][] {
+   *         {-14.276086, -12.433481, -2.2447076, -1.5775859, 1.8588694}
+   *     });
+   * Operand<TFloat64> result = fops.linalg.matmul(a, b);
+   * // result = {
+   * //     {127.69746,  111.21564,  20.078575,  14.111271,  -16.62731},
+   * //     {-59.547394, -51.861652, -9.362965,  -6.580314,    7.753584},
+   * //     {  3.1950197,  2.7826407, 0.50237054, 0.35306725, -0.4160191}
+   * //  }
+   *
+   * }</pre>
+   *
+   * @param scope The TensorFlow scope
+   * @param a an Operand of of type {@code TFloat16}, {@code TFloat32}, {@code TFloat64 }, {@code
+   *     TInt32}. with a {@code rank > 1}
+   * @param b an Operand with same type and rank as {@code a}.
+   * @param transposeA If true, {@code a} is transposed before multiplication.
+   * @param transposeB If true, {@code b} is transposed before multiplication
+   * @param <T> the data type of the Operands
+   * @return A Operand of the same type as {@code a} and {@code b} where each inner-most matrix is
+   *     the product of the corresponding matrices in {@code a} and {@code b}. This is the matrix
+   *     product not an element-wise product.
+   * @throws java.lang.IllegalArgumentException If {@code transposeA} and {@code adjointA} , or
+   *     {@code transposeB} and {@code adjointB} are both set to `true`.
+   */
+  public static <T extends TNumber> Operand<T> matmul(
+      Scope scope, Operand<T> a, Operand<T> b, boolean transposeA, boolean transposeB) {
+    return matmul(scope, a, b, transposeA, transposeB, false, false, false, false);
+  }
+
+  /**
+   * Multiplies matrix {@code a} by matrix {@code b}, producing {@code a} * {@code b }.
+   *
+   * <p>The inputs must, following any transpositions, be tensors of {@code rank >= 2} where the
+   * inner 2 dimensions specify valid matrix multiplication dimensions, and any further outer
+   * dimensions specify matching batch size.
+   *
+   * <p>Both matrices must be of the same type. The supported types are: {@code TFloat16}, {@code
+   * TFloat32}, {@code TFloat64}, {@code TInt32}.
+   *
+   * <p>Either matrix can be transposed or adjointed (conjugated and transposed) on the fly by
+   * setting one of the corresponding flag to true. These are false by default.
+   *
+   * <p>Note: This is matrix product, not element-wise product.
+   *
+   * <p>A simple 2-D tensor matrix multiplication:
+   *
+   * <pre>{@code
+   * Operand<TFloat64> a = tf.constant(new double[][] {
+   *         {-8.944851},
+   *         {4.1711287},
+   *         {-0.22380222}
+   *     });
+   * Operand<TFloat64> b = tf.constant( new double[][] {
+   *         {-14.276086, -12.433481, -2.2447076, -1.5775859, 1.8588694}
+   *     });
+   * Operand<TFloat64> result = fops.linalg.matmul(a, b);
+   * // result = {
+   * //     {127.69746,  111.21564,  20.078575,  14.111271,  -16.62731},
+   * //     {-59.547394, -51.861652, -9.362965,  -6.580314,    7.753584},
+   * //     {  3.1950197,  2.7826407, 0.50237054, 0.35306725, -0.4160191}
+   * //  }
+   *
+   * }</pre>
+   *
+   * @param scope The TensorFlow scope
+   * @param a an Operand of of type {@code TFloat16}, {@code TFloat32}, {@code TFloat64 }, {@code
+   *     TInt32}. with a {@code rank > 1}
+   * @param b an Operand with same type and rank as {@code a}.
+   * @param transposeA If true, {@code a} is transposed before multiplication.
+   * @param transposeB If True, {@code b} is transposed before multiplication
+   * @param adjointA If true, {@code a} is conjugated and transposed before multiplication.
+   * @param adjointB If true, {@code b} is conjugated and transposed before multiplication.
+   * @param aIsSparse If true, {@code a} is treated as a sparse matrix. Notice, this <em>does not
+   *     support {@link SparseTensor}</em>, it just makes optimizations that assume most values in
+   *     {@code a} are zero.
+   * @param bIsSparse If true, {@code b} is treated as a sparse matrix. Notice, this <em>does not
+   *     support {@link SparseTensor}</em>, it just makes optimizations that assume most values in
+   *     {@code b} are zero.
+   * @param <T> the data type of the Operands
+   * @return A Operand of the same type as {@code a} and {@code b} where each inner-most matrix is
+   *     the product of the corresponding matrices in {@code a} and {@code b}. This is the matrix
+   *     product not an element-wise product.
+   * @throws java.lang.IllegalArgumentException If {@code transposeA} and {@code adjointA} , or
+   *     {@code transposeB} and {@code adjointB} are both set to `true`.
+   */
+  @SuppressWarnings("unchecked")
+  public static <T extends TNumber> Operand<T> matmul(
+      Scope scope,
+      Operand<T> a,
+      Operand<T> b,
+      boolean transposeA,
+      boolean transposeB,
+      boolean adjointA,
+      boolean adjointB,
+      boolean aIsSparse,
+      boolean bIsSparse) {
+    Scope lscope = scope.withSubScope("MatMul");
+    if (transposeA && adjointA)
+      throw new IllegalArgumentException("Only one of transposeA and adjointA can be true.");
+    if (transposeB && adjointB)
+      throw new IllegalArgumentException("Only one of transposeB and adjointB can be true.");
+    if (!(TFloating.class.isAssignableFrom(a.type()) || a.type().equals(TInt32.class)))
+      throw new IllegalArgumentException(
+          String.format(
+              "Operand 'a' must be of type 'TBfloat16','TFloat16', 'TFloat32', 'TFloat64' or 'TInt32'. found type : %s",
+              a.type().getSimpleName()));
+    if (!(TFloating.class.isAssignableFrom(a.type()) || b.type().equals(TInt32.class)))
+      throw new IllegalArgumentException(
+          String.format(
+              "Operand 'b' must be of type 'TBfloat16', 'TFloat32', 'TFloat64' or 'TInt32'. found type : %s",
+              b.type().getSimpleName()));
+
+    Shape aShape = a.shape();
+    Shape bShape = b.shape();
+    if (aShape.numDimensions() != bShape.numDimensions())
+      throw new IllegalArgumentException(
+          String.format(
+              "Parameters 'a' and 'b' must the same rank: found a rank = %d, b rank = %d",
+              aShape.numDimensions(), bShape.numDimensions()));
+    boolean outputMayHaveNonEmptyBatchShape =
+        aShape.numDimensions() == Shape.UNKNOWN_SIZE
+            || aShape.numDimensions() > 2
+            || bShape.numDimensions() == Shape.UNKNOWN_SIZE;
+
+    if ((!aIsSparse && !bIsSparse) && outputMayHaveNonEmptyBatchShape) {
+      // BatchMatmul does not support transpose, so we conjugate the matrix and
+      // use adjoint instead. Conj() is a noop for real matrices.
+      if (transposeA) {
+        a = Conj.create(scope, a);
+        adjointA = true;
+      }
+      if (transposeB) {
+        b = Conj.create(scope, b);
+        adjointB = true;
+      }
+      return BatchMatMul.create(
+          lscope, a, b, BatchMatMul.adjX(adjointA), BatchMatMul.adjY(adjointB));
+    }
+
+    // Neither matmul nor sparse_matmul support adjoint, so we conjugate
+    // the matrix and use transpose instead. Conj() is a noop for real
+    // matrices.
+    if (adjointA) {
+      a = Conj.create(scope, a);
+      transposeA = true;
+    }
+    if (adjointB) {
+      b = Conj.create(scope, b);
+      transposeB = true;
+    }
+
+    boolean useSparseMatmul = false;
+    if (aIsSparse || bIsSparse) {
+      useSparseMatmul =
+          (a.type().equals(TBfloat16.class) || a.type().equals(TFloat32.class))
+              && (b.type().equals(TBfloat16.class) || b.type().equals(TFloat32.class));
+    }
+    if ((a.type().equals(TBfloat16.class) || b.type().equals(TBfloat16.class))
+        && !a.type().equals(b.type())) useSparseMatmul = true;
+
+    if (useSparseMatmul) {
+      Operand<TFloat32> result =
+          SparseMatMul.create(
+              lscope,
+              a,
+              b,
+              SparseMatMul.transposeA(transposeA),
+              SparseMatMul.transposeB(transposeB),
+              SparseMatMul.aIsSparse(aIsSparse),
+              SparseMatMul.bIsSparse(bIsSparse));
+      if (a.type().equals(TFloat32.class)) return (Operand<T>) result;
+      else return Cast.create(scope, result, a.type());
+    }
+
+    return org.tensorflow.op.linalg.MatMul.create(
+        lscope,
+        a,
+        b,
+        org.tensorflow.op.linalg.MatMul.transposeA(transposeA),
+        org.tensorflow.op.linalg.MatMul.transposeB(transposeB));
+  }
+}
diff --git a/tensorflow-framework/src/main/java/org/tensorflow/framework/op/math/Axes.java b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/math/Axes.java
new file mode 100644
index 00000000000..cea7c5f80bf
--- /dev/null
+++ b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/math/Axes.java
@@ -0,0 +1,49 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+=======================================================================*/
+package org.tensorflow.framework.op.math;
+
+import org.tensorflow.Operand;
+import org.tensorflow.ndarray.Shape;
+import org.tensorflow.op.Scope;
+import org.tensorflow.op.core.Constant;
+import org.tensorflow.op.core.Range;
+import org.tensorflow.op.core.Rank;
+import org.tensorflow.types.TInt32;
+import org.tensorflow.types.family.TType;
+
+/** Axes Operations */
+public class Axes {
+
+  /**
+   * Creates an Operand that has all axes contained in the Operand's shape.
+   *
+   * @param scope The TensorFlow scope
+   * @param op the Operand
+   * @return an Operand that has all axes contained in the Operand's shape..
+   */
+  public static Operand<TInt32> allAxes(Scope scope, Operand<? extends TType> op) {
+    int rank = op.shape().numDimensions();
+    if (rank != Shape.UNKNOWN_SIZE) {
+      int[] axes = new int[rank];
+      for (int i = 0; i < rank; i++) {
+        axes[i] = i;
+      }
+      return Constant.vectorOf(scope, axes);
+    } else {
+      return Range.create(
+          scope, Constant.scalarOf(scope, 0), Rank.create(scope, op), Constant.scalarOf(scope, 1));
+    }
+  }
+}
diff --git a/tensorflow-framework/src/main/java/org/tensorflow/framework/op/math/ConfusionMatrix.java b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/math/ConfusionMatrix.java
new file mode 100644
index 00000000000..71103133066
--- /dev/null
+++ b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/math/ConfusionMatrix.java
@@ -0,0 +1,318 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+=======================================================================*/
+package org.tensorflow.framework.op.math;
+
+import java.util.Arrays;
+import java.util.Collections;
+import org.tensorflow.Operand;
+import org.tensorflow.framework.losses.impl.LossTuple;
+import org.tensorflow.ndarray.Shape;
+import org.tensorflow.op.Scope;
+import org.tensorflow.op.core.AssertThat;
+import org.tensorflow.op.core.Constant;
+import org.tensorflow.op.core.Identity;
+import org.tensorflow.op.core.OnesLike;
+import org.tensorflow.op.core.ReduceAll;
+import org.tensorflow.op.core.ReduceMax;
+import org.tensorflow.op.core.ScatterNd;
+import org.tensorflow.op.core.Squeeze;
+import org.tensorflow.op.core.Stack;
+import org.tensorflow.op.dtypes.Cast;
+import org.tensorflow.op.math.Add;
+import org.tensorflow.op.math.GreaterEqual;
+import org.tensorflow.op.math.Less;
+import org.tensorflow.op.math.Maximum;
+import org.tensorflow.types.TBool;
+import org.tensorflow.types.TInt64;
+import org.tensorflow.types.family.TNumber;
+
+/** Confusion Matrix Operations */
+public class ConfusionMatrix {
+
+  /**
+   * Computes the confusion matrix from predictions and labels.
+   *
+   * <p>The matrix columns represent the prediction labels and the rows represent the real labels.
+   * The confusion matrix is always a 2-D array of shape {@code [n,n]}, where {@code n} is the
+   * number of valid labels for a given classification task. Both prediction and labels must be 1-D
+   * arrays of the same shape in order for this function to work.
+   *
+   * <p>If {@code numClasses} is null, then {@code numClasses} will be set to one plus the maximum
+   * value in either predictions or labels. Class labels are expected to start at 0. For example, if
+   * {@code numClasses} is 3, then the possible labels would be {@code [0, 1, 2]}.
+   *
+   * <p>If {@code weights} is not null, then each prediction contributes its corresponding weight to
+   * the total value of the confusion matrix cell.
+   *
+   * <p>For example:
+   *
+   * <pre>{@code
+   * fops.math.confusion_matrix(tf.constant(new int[] {1, 2, 4}), tf.constant(new int[] {2, 2, 4})) ==>
+   *     [[0 0 0 0 0]
+   *      [0 0 1 0 0]
+   *      [0 0 1 0 0]
+   *      [0 0 0 0 0]
+   *      [0 0 0 0 1]]
+   * }</pre>
+   *
+   * <p>Note that the possible labels are assumed to be {@code [0, 1, 2, 3, 4]}, resulting in a 5x5
+   * confusion matrix.
+   *
+   * @param scope The TensorFlow scope
+   * @param labels 1-D Operand of real labels for the classification task.
+   * @param predictions 1-D Operand of predictions for a given classification.
+   * @param <T> Data type of the confusion matrix.
+   * @return An Operand of type {@code type} with shape {@code [n, n]} representing the confusion
+   *     matrix, where {@code n} is the number of possible labels in the classification task.
+   * @throws IllegalArgumentException If both predictions and labels are not 1-D vectors and have
+   *     mismatched shapes, or if {@code weights} is not null and its shape doesn't match {@code
+   *     predictions}.
+   */
+  public static <T extends TNumber> Operand<T> confusionMatrix(
+      Scope scope, Operand<T> labels, Operand<T> predictions) {
+    return confusionMatrix(scope, labels, predictions, null, null);
+  }
+
+  /**
+   * Computes the confusion matrix from predictions and labels.
+   *
+   * <p>The matrix columns represent the prediction labels and the rows represent the real labels.
+   * The confusion matrix is always a 2-D array of shape {@code [n,n]}, where {@code n} is the
+   * number of valid labels for a given classification task. Both prediction and labels must be 1-D
+   * arrays of the same shape in order for this function to work.
+   *
+   * <p>If {@code numClasses} is null, then {@code numClasses} will be set to one plus the maximum
+   * value in either predictions or labels. Class labels are expected to start at 0. For example, if
+   * {@code numClasses} is 3, then the possible labels would be {@code [0, 1, 2]}.
+   *
+   * <p>If {@code weights} is not null, then each prediction contributes its corresponding weight to
+   * the total value of the confusion matrix cell.
+   *
+   * <p>For example:
+   *
+   * <pre>{@code
+   * fops.math.confusion_matrix(tf.constant(new int[] {1, 2, 4}), tf.constant(new int[] {2, 2, 4})) ==>
+   *     [[0 0 0 0 0]
+   *      [0 0 1 0 0]
+   *      [0 0 1 0 0]
+   *      [0 0 0 0 0]
+   *      [0 0 0 0 1]]
+   * }</pre>
+   *
+   * <p>Note that the possible labels are assumed to be {@code [0, 1, 2, 3, 4]}, resulting in a 5x5
+   * confusion matrix.
+   *
+   * @param scope The TensorFlow scope
+   * @param labels 1-D Operand of real labels for the classification task.
+   * @param predictions 1-D Operand of predictions for a given classification.
+   * @param weights An optional Operand whose shape matches {@code predictions}.
+   * @param <T> Data type of the confusion matrix.
+   * @return An Operand of type {@code type} with shape {@code [n, n]} representing the confusion
+   *     matrix, where {@code n} is the number of possible labels in the classification task.
+   * @throws IllegalArgumentException If both predictions and labels are not 1-D vectors and have
+   *     mismatched shapes, or if {@code weights} is not null and its shape doesn't match {@code
+   *     predictions}.
+   */
+  public static <T extends TNumber> Operand<T> confusionMatrix(
+      Scope scope, Operand<T> labels, Operand<T> predictions, Operand<T> weights) {
+    return confusionMatrix(scope, labels, predictions, weights, null);
+  }
+
+  /**
+   * Computes the confusion matrix from predictions and labels.
+   *
+   * <p>The matrix columns represent the prediction labels and the rows represent the real labels.
+   * The confusion matrix is always a 2-D array of shape {@code [n,n]}, where {@code n} is the
+   * number of valid labels for a given classification task. Both prediction and labels must be 1-D
+   * arrays of the same shape in order for this function to work.
+   *
+   * <p>If {@code numClasses} is null, then {@code numClasses} will be set to one plus the maximum
+   * value in either predictions or labels. Class labels are expected to start at 0. For example, if
+   * {@code numClasses} is 3, then the possible labels would be {@code [0, 1, 2]}.
+   *
+   * <p>If {@code weights} is not null, then each prediction contributes its corresponding weight to
+   * the total value of the confusion matrix cell.
+   *
+   * <p>For example:
+   *
+   * <pre>{@code
+   * fops.math.confusion_matrix(tf.constant(new int[] {1, 2, 4}), tf.constant(new int[] {2, 2, 4})) ==>
+   *     [[0 0 0 0 0]
+   *      [0 0 1 0 0]
+   *      [0 0 1 0 0]
+   *      [0 0 0 0 0]
+   *      [0 0 0 0 1]]
+   * }</pre>
+   *
+   * <p>Note that the possible labels are assumed to be {@code [0, 1, 2, 3, 4]}, resulting in a 5x5
+   * confusion matrix.
+   *
+   * @param scope The TensorFlow scope
+   * @param labels 1-D Operand of real labels for the classification task.
+   * @param predictions 1-D Operand of predictions for a given classification.
+   * @param weights An optional Operand whose shape matches {@code predictions}.
+   * @param numClasses The possible number of labels the classification task can have. If this value
+   *     is null, it will be calculated using both predictions and labels.
+   * @param <T> Data type of the confusion matrix.
+   * @return An Operand of type {@code type} with shape {@code [n, n]} representing the confusion
+   *     matrix, where {@code n} is the number of possible labels in the classification task.
+   * @throws IllegalArgumentException If both predictions and labels are not 1-D vectors and have
+   *     mismatched shapes, or if {@code weights} is not null and its shape doesn't match {@code
+   *     predictions}.
+   */
+  public static <T extends TNumber> Operand<T> confusionMatrix(
+      Scope scope,
+      Operand<T> labels,
+      Operand<T> predictions,
+      Operand<T> weights,
+      Operand<TInt64> numClasses) {
+    Scope lScope = scope.withSubScope("confusionMatrix");
+    LossTuple<T> tuple = removeSqueezableDimensions(scope, labels, predictions, 0);
+    Operand<TInt64> lLabels = Cast.create(lScope, tuple.getLabels(), TInt64.class);
+    Operand<TInt64> lPredictions = Cast.create(lScope, tuple.getTarget(), TInt64.class);
+
+    Operand<TInt64> zero = Constant.scalarOf(lScope, 0L);
+    Operand<TInt64> one = Constant.scalarOf(lScope, 1L);
+
+    AssertThat labelsNonNegative =
+        AssertThat.create(
+            lScope,
+            ReduceAll.create(
+                lScope, GreaterEqual.create(lScope, lLabels, zero), Axes.allAxes(scope, lLabels)),
+            Collections.singletonList(
+                Constant.scalarOf(lScope, "labels contains negative values")));
+    lLabels =
+        Identity.create(
+            lScope.withControlDependencies(Collections.singletonList(labelsNonNegative)), lLabels);
+
+    AssertThat predictionsNonNegative =
+        AssertThat.create(
+            lScope,
+            ReduceAll.create(
+                lScope,
+                GreaterEqual.create(lScope, lPredictions, zero),
+                Axes.allAxes(scope, lPredictions)),
+            Collections.singletonList(
+                Constant.scalarOf(lScope, "predictions contains negative values")));
+    lPredictions =
+        Identity.create(
+            lScope.withControlDependencies(Collections.singletonList(predictionsNonNegative)),
+            lPredictions);
+
+    Operand<TInt64> lNumClasses;
+    if (numClasses == null) {
+      lNumClasses =
+          Add.create(
+              lScope,
+              Maximum.create(
+                  lScope,
+                  ReduceMax.create(lScope, lPredictions, zero),
+                  ReduceMax.create(lScope, lLabels, zero)),
+              one);
+    } else {
+      lNumClasses = Cast.create(lScope, numClasses, TInt64.class);
+      Operand<TBool> less = Less.create(lScope, lLabels, lNumClasses);
+      AssertThat labelsLess =
+          AssertThat.create(
+              lScope,
+              ReduceAll.create(scope, less, Axes.allAxes(scope, less), ReduceAll.keepDims(false)),
+              Collections.singletonList(Constant.scalarOf(lScope, "labels out of bounds")));
+      lLabels =
+          Identity.create(
+              lScope.withControlDependencies(Collections.singletonList(labelsLess)), lLabels);
+
+      less = Less.create(lScope, lPredictions, lNumClasses);
+      AssertThat predictionsLess =
+          AssertThat.create(
+              lScope,
+              ReduceAll.create(scope, less, Axes.allAxes(scope, less), ReduceAll.keepDims(false)),
+              Collections.singletonList(Constant.scalarOf(lScope, "predictions  out of bounds")));
+      lPredictions =
+          Identity.create(
+              lScope.withControlDependencies(Collections.singletonList(predictionsLess)),
+              lPredictions);
+    }
+
+    if (weights != null) {
+      if (!predictions.shape().isCompatibleWith(weights.shape())) {
+        throw new IllegalArgumentException(
+            String.format(
+                "predictions.shape() [%s], is not compatible with weights.shape() [ %s].",
+                predictions.shape(), weights.shape()));
+      }
+    }
+
+    Operand<TInt64> shape = Stack.create(lScope, Arrays.asList(lNumClasses, lNumClasses));
+    Operand<TInt64> indices =
+        Stack.create(lScope, Arrays.asList(lLabels, lPredictions), Stack.axis(1L));
+    Operand<T> values = weights == null ? OnesLike.create(lScope, predictions) : weights;
+    /// Operand<T> zeroMatrix = Zeros.create(lScope, Cast.create(lScope, shape, TInt32.class),
+    // type);
+
+    return ScatterNd.create(lScope, indices, values, shape);
+  }
+
+  /**
+   * Squeeze last dim if ranks differ from expected by exactly 1.
+   *
+   * @param scope The TensorFlow scope
+   * @param labels Label values, a {@code Operand} whose dimensions match {@code predictions }.
+   * @param predictions Predicted values, a {@code Tensor} of arbitrary dimensions.
+   * @param expectedRankDiff Expected result of {@code rank(predictions) - rank(labels)}.
+   * @param <T> the data type for the labels, predictions and result
+   * @return {@code labels} and {@code predictions}, possibly with last dim squeezed.
+   */
+  private static <T extends TNumber> LossTuple<T> removeSqueezableDimensions(
+      Scope scope, Operand<T> labels, Operand<T> predictions, int expectedRankDiff) {
+    Scope lScope = scope.withSubScope("removeSqueezableDimensions");
+    Shape predictionsShape = predictions.shape();
+    int predictionsRank = predictionsShape.numDimensions();
+    Shape labelsShape = labels.shape();
+    int labelsRank = labelsShape.numDimensions();
+
+    if (predictionsRank != Shape.UNKNOWN_SIZE || labelsRank != Shape.UNKNOWN_SIZE) {
+      // Use rank.
+      int rankDiff = predictionsRank - labelsRank;
+      if (rankDiff == expectedRankDiff + 1 && Shape.isCompatible(predictionsShape.size(-1), 1)) {
+        predictions = Squeeze.create(lScope, predictions);
+      } else if (rankDiff == expectedRankDiff - 1 && Shape.isCompatible(labelsShape.size(-1), 1)) {
+        labels = Squeeze.create(lScope, labels);
+      }
+      return new LossTuple<>(labels, predictions);
+    }
+    // Use dynamic rank.
+
+    // TODO: hold for lazy select feature,
+    //  Operand<TInt32> rankDiff = tf.math.sub(tf.rank(predictions), tf.rank(labels));
+    if (predictionsRank == Shape.UNKNOWN_SIZE && Shape.isCompatible(predictionsShape.size(-1), 1)) {
+      /*
+       * TODO, if we ever get a select that does lazy evaluation, but for now do the tf.squeeze
+       * predictions = tf.select( tf.math.equal(tf.constant(expectedRankDiff+1),rankDiff ),
+       * tf.squeeze(predictions, Squeeze.axis(Arrays.asList(-1L))), predictions ); *
+       */
+      predictions =
+          Squeeze.create(lScope, predictions, Squeeze.axis(Collections.singletonList(-1L)));
+    }
+    if (labelsRank == Shape.UNKNOWN_SIZE && Shape.isCompatible(labelsShape.size(-1), 1)) {
+      /*
+       * TODO, if we ever get a select that does lazy evaluation labels = tf.select(
+       * tf.math.equal(tf.constant(expectedRankDiff+1),rankDiff ), tf.squeeze(labels,
+       * Squeeze.axis(Arrays.asList(-1L))), predictions ); *
+       */
+      labels = Squeeze.create(lScope, labels, Squeeze.axis(Collections.singletonList(-1L)));
+    }
+    return new LossTuple<>(labels, predictions);
+  }
+}
diff --git a/tensorflow-framework/src/main/java/org/tensorflow/framework/op/math/L2Normalize.java b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/math/L2Normalize.java
new file mode 100644
index 00000000000..f8cdfe29026
--- /dev/null
+++ b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/math/L2Normalize.java
@@ -0,0 +1,54 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+=======================================================================*/
+package org.tensorflow.framework.op.math;
+
+import org.tensorflow.Operand;
+import org.tensorflow.op.Scope;
+import org.tensorflow.op.core.Constant;
+import org.tensorflow.op.core.ReduceSum;
+import org.tensorflow.op.dtypes.Cast;
+import org.tensorflow.op.math.Maximum;
+import org.tensorflow.op.math.Mul;
+import org.tensorflow.op.math.Rsqrt;
+import org.tensorflow.op.math.Square;
+import org.tensorflow.types.family.TNumber;
+
+/** L2 Normalization Operations */
+public class L2Normalize {
+
+  /**
+   * Normalizes along dimension axis using an L2 norm.
+   *
+   * @param scope The TensorFlow scope
+   * @param x the input
+   * @param axis Dimension along which to normalize.
+   * @param <T> the data type for the input and the result
+   * @return the normalized values based on L2 norm
+   */
+  public static <T extends TNumber> Operand<T> l2Normalize(Scope scope, Operand<T> x, int[] axis) {
+    Operand<T> squareSum =
+        ReduceSum.create(
+            scope,
+            Square.create(scope, x),
+            Constant.vectorOf(scope, axis),
+            ReduceSum.keepDims(Boolean.TRUE));
+    Operand<T> invNorm =
+        Rsqrt.create(
+            scope,
+            Maximum.create(
+                scope, squareSum, Cast.create(scope, Constant.scalarOf(scope, 1e-12F), x.type())));
+    return Mul.create(scope, x, invNorm);
+  }
+}
diff --git a/tensorflow-framework/src/main/java/org/tensorflow/framework/op/math/ReduceLogSumExp.java b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/math/ReduceLogSumExp.java
new file mode 100644
index 00000000000..83660befe3a
--- /dev/null
+++ b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/math/ReduceLogSumExp.java
@@ -0,0 +1,171 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+=======================================================================*/
+package org.tensorflow.framework.op.math;
+
+import org.tensorflow.Operand;
+import org.tensorflow.ndarray.Shape;
+import org.tensorflow.op.Scope;
+import org.tensorflow.op.core.Constant;
+import org.tensorflow.op.core.Range;
+import org.tensorflow.op.core.Rank;
+import org.tensorflow.op.core.ReduceMax;
+import org.tensorflow.op.core.ReduceSum;
+import org.tensorflow.op.core.Reshape;
+import org.tensorflow.op.core.Select;
+import org.tensorflow.op.core.StopGradient;
+import org.tensorflow.op.core.ZerosLike;
+import org.tensorflow.op.math.Add;
+import org.tensorflow.op.math.Exp;
+import org.tensorflow.op.math.IsFinite;
+import org.tensorflow.op.math.Log;
+import org.tensorflow.op.math.Sub;
+import org.tensorflow.types.TInt32;
+import org.tensorflow.types.family.TFloating;
+
+/** Reduce Log Sum Exp Operations */
+public class ReduceLogSumExp {
+
+  /**
+   * Computes log(sum(exp(elements across dimensions of a tensor))). Reduces {@code input_tensor}
+   * along the dimensions given in {@code axes}.
+   *
+   * <p>Reduces {@code input} along the dimensions given in {@code axes}. Unless {@code keepdims} is
+   * true, the rank of the tensor is reduced by 1 for each of the entries in {@code axes}, which
+   * must be unique. If {@code keepdims} is true, the reduced dimensions are retained with length 1.
+   * If {@code axes} has no entries, all dimensions are reduced, and a tensor with a single element
+   * is returned. This function is more numerically stable than {@code log(sum(exp(input)))}. It
+   * avoids overflows caused by taking the exp of large inputs and underflows caused by taking the
+   * log of small inputs.
+   *
+   * @param scope The TensorFlow scope
+   * @param input The tensor to reduce.
+   * @param axes The dimensions to reduce. If null, reduces all dimensions. Must be in the range
+   *     {@code [-rank(input_tensor), rank(input_tensor)]}.
+   * @param keepDims If true, retains reduced dimensions with length 1.
+   * @param <T> the data type for the input and the result
+   * @return The reduced tensor.
+   */
+  public static <T extends TFloating> Operand<T> reduceLogSumExp(
+      Scope scope, Operand<T> input, int[] axes, boolean keepDims) {
+    Operand<TInt32> reduceDims = reductionDims(scope, input, axes);
+    Operand<T> rawMax = reduceMaxWithDims(scope, input, axes, keepDims, reduceDims);
+    Operand<T> myMax =
+        StopGradient.create(
+            scope,
+            Select.create(
+                scope, IsFinite.create(scope, rawMax), rawMax, ZerosLike.create(scope, rawMax)));
+
+    Operand<T> result =
+        Log.create(
+            scope,
+            reduceSumWithDims(
+                scope,
+                Exp.create(scope, Sub.create(scope, input, myMax)),
+                axes,
+                keepDims,
+                reduceDims));
+
+    if (!keepDims) {
+      myMax = Reshape.create(scope, myMax, org.tensorflow.op.core.Shape.create(scope, result));
+    }
+    result = Add.create(scope, result, myMax);
+    return mayReduceToScalar(scope, keepDims, axes, result);
+  }
+
+  /**
+   * @param scope The TensorFlow scope
+   * @param input The tensor to reduce.
+   * @param axes The dimensions to reduce. If null, reduces all dimensions. Must be in the range
+   *     {@code [-rank(input_tensor), rank(input_tensor)]}.
+   * @param keepDims If true, retains reduced dimensions with length 1.
+   * @param dims the reduction dimensions
+   * @param <T> the data type for the input and the result
+   * @return the reduced sum
+   */
+  private static <T extends TFloating> Operand<T> reduceSumWithDims(
+      Scope scope, Operand<T> input, int[] axes, boolean keepDims, Operand<TInt32> dims) {
+    return mayReduceToScalar(
+        scope, keepDims, axes, ReduceSum.create(scope, input, dims, ReduceSum.keepDims(keepDims)));
+  }
+
+  /**
+   * @param scope The TensorFlow scope
+   * @param input The tensor to reduce.
+   * @param axes The dimensions to reduce. If null, reduces all dimensions. Must be in the range
+   *     {@code [-rank(input_tensor), rank(input_tensor)]}.
+   * @param keepDims If true, retains reduced dimensions with length 1.
+   * @param dims the reduction dimensions
+   * @param <T> the data type for the input and the result
+   * @return the reduced maximum input value
+   */
+  private static <T extends TFloating> Operand<T> reduceMaxWithDims(
+      Scope scope, Operand<T> input, int[] axes, boolean keepDims, Operand<TInt32> dims) {
+    return mayReduceToScalar(
+        scope, keepDims, axes, ReduceMax.create(scope, input, dims, ReduceMax.keepDims(keepDims)));
+  }
+
+  /**
+   * Sets a reduction's output shape to be a scalar if possible.
+   *
+   * @param scope The TensorFlow scope
+   * @param keepDims If true, retains reduced dimensions with length 1.
+   * @param axes The dimensions to reduce. If null, reduces all dimensions. Must be in the range
+   *     {@code [-rank(input_tensor), rank(input_tensor)]}.
+   * @param output the output, possibly reduced to a scalar
+   * @param <T> the datat type of the Operands.
+   * @return the operand, possibly reduced to a scalar.
+   */
+  private static <T extends TFloating> Operand<T> mayReduceToScalar(
+      Scope scope, boolean keepDims, int[] axes, Operand<T> output) {
+
+    if ((output.shape().numDimensions() == Shape.UNKNOWN_SIZE
+            || output.shape().hasUnknownDimension())
+        && !keepDims
+        && axes == null) {
+      return Reshape.create(scope, output, Constant.tensorOf(scope, Shape.scalar()));
+    } else {
+      return output;
+    }
+  }
+
+  /**
+   * Reduce dimensions based on axis
+   *
+   * @param input the input
+   * @param axes he dimensions to reduce, may be null
+   * @return the dimensions to be reduced.
+   */
+  private static <T extends TFloating> Operand<TInt32> reductionDims(
+      Scope scope, Operand<T> input, int[] axes) {
+    if (axes != null) {
+      return Constant.vectorOf(scope, axes);
+    }
+    long rank = input.shape().numDimensions();
+    if (rank != Shape.UNKNOWN_SIZE) {
+      int[] dims = new int[(int) rank];
+      for (int i = 0; i < rank; i++) {
+        dims[i] = i;
+      }
+      return Constant.vectorOf(scope, dims);
+
+    } else {
+      return Range.create(
+          scope,
+          Constant.scalarOf(scope, 0),
+          Rank.create(scope, input),
+          Constant.scalarOf(scope, 1));
+    }
+  }
+}
diff --git a/tensorflow-framework/src/main/java/org/tensorflow/framework/op/math/TensorDot.java b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/math/TensorDot.java
new file mode 100644
index 00000000000..a222f64679e
--- /dev/null
+++ b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/math/TensorDot.java
@@ -0,0 +1,663 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+=======================================================================*/
+package org.tensorflow.framework.op.math;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.tensorflow.Graph;
+import org.tensorflow.Operand;
+import org.tensorflow.Session;
+import org.tensorflow.framework.op.linalg.MatMul;
+import org.tensorflow.ndarray.Shape;
+import org.tensorflow.op.Scope;
+import org.tensorflow.op.core.AssertThat;
+import org.tensorflow.op.core.Concat;
+import org.tensorflow.op.core.Constant;
+import org.tensorflow.op.core.Gather;
+import org.tensorflow.op.core.Range;
+import org.tensorflow.op.core.Rank;
+import org.tensorflow.op.core.ReduceProd;
+import org.tensorflow.op.core.Reshape;
+import org.tensorflow.op.core.Select;
+import org.tensorflow.op.core.SetDiff1d;
+import org.tensorflow.op.core.Slice;
+import org.tensorflow.op.core.Stack;
+import org.tensorflow.op.dtypes.Cast;
+import org.tensorflow.op.linalg.Transpose;
+import org.tensorflow.op.math.Add;
+import org.tensorflow.op.math.GreaterEqual;
+import org.tensorflow.op.math.Less;
+import org.tensorflow.op.math.Sub;
+import org.tensorflow.types.TBfloat16;
+import org.tensorflow.types.TFloat16;
+import org.tensorflow.types.TInt32;
+import org.tensorflow.types.family.TFloating;
+import org.tensorflow.types.family.TNumber;
+
+/** tensor contraction Operations */
+public class TensorDot {
+
+  /**
+   * Transpose and reshape the input for contraction op.
+   *
+   * <p>This method is helpful in reducing {@code math.tensordot} to {@code math_ops.matmul} using
+   * {@code array_ops.transpose} and {@code array_ops.reshape}. The method takes a tensor and
+   * performs the correct transpose and reshape operation for a given set of indices. It returns the
+   * reshaped tensor as well as a list of indices necessary to reshape the tensor again after matrix
+   * multiplication.
+   *
+   * @param <T> the type of Operand
+   * @param a the Tensor
+   * @param axis unique indices specifying valid axes of {@code a}.
+   * @param flipped whether to flip the dimensions or not
+   * @return A tuple (reshapedA, freeDims, freeDimsStatic) where reshapedA is a reshaped to allow
+   *     contraction via matmul, freeDims is a TInt32 Operand, depending on whether the shape of a
+   *     is fully specified, and freeDimsStatic is either a list of integers and null values, or
+   *     None, representing the inferred shape of the free dimensions
+   */
+  private static <T extends TNumber> Object[] tensordotReshape(
+      Scope scope, Operand<T> a, Operand<TInt32> axis, boolean flipped) {
+    Shape aShape = a.shape();
+
+    if (!aShape.hasUnknownDimension()) { // calculate using values
+      long[] aShapeDims = aShape.asArray();
+      if (aShapeDims == null) aShapeDims = new long[0];
+      long[] aDimsIndex = new long[aShapeDims.length];
+      for (int i = 0; i < aDimsIndex.length; i++) aDimsIndex[i] = i;
+
+      // get int array from axis Operand
+      int[] iAxes = getIntArray(scope, axis);
+      // Convert negative axes to positive
+      for (int i = 0; i < iAxes.length; i++)
+        iAxes[i] = iAxes[i] >= 0 ? iAxes[i] : Math.floorMod(iAxes[i], iAxes.length);
+
+      // convert integer axis to long axis
+      long[] lAxes = Arrays.stream(iAxes).mapToLong(i -> i).toArray();
+
+      // create list of the axes, dims, and free axes
+      List<Long> axesList = Arrays.stream(lAxes).boxed().collect(Collectors.toList());
+      List<Long> freeList = Arrays.stream(aDimsIndex).boxed().collect(Collectors.toList());
+      freeList.removeAll(axesList);
+
+      // create array of free dims
+      long[] free = freeList.stream().mapToLong(i -> i).toArray();
+      long[] freeDims = new long[free.length];
+      for (int i = 0; i < free.length; i++) freeDims[i] = aShapeDims[(int) free[i]];
+
+      // Calculate the free dim by doing a reduce prod
+      long prodFree = 1;
+      for (long i : freeDims) {
+        prodFree *= i;
+      }
+
+      // calculate the used dims by doing a reduce prod
+      long prodAxis = 1;
+      for (long i : lAxes) {
+        prodAxis *= aShapeDims[(int) i];
+      }
+
+      // setup the permutations array for the transpose
+      long[] perm = new long[freeDims.length + lAxes.length];
+      Shape newShape;
+      if (flipped) {
+        System.arraycopy(lAxes, 0, perm, 0, lAxes.length);
+        System.arraycopy(free, 0, perm, lAxes.length, free.length);
+        newShape = Shape.of(prodAxis, prodFree);
+      } else {
+        System.arraycopy(free, 0, perm, 0, free.length);
+        System.arraycopy(lAxes, 0, perm, freeDims.length, lAxes.length);
+        newShape = Shape.of(prodFree, prodAxis);
+      }
+
+      Operand<T> aTrans;
+      long[] arrange = new long[lAxes.length];
+      for (int i = 0; i < arrange.length; i++) arrange[i] = i;
+
+      // if the permutations is not equals to the natural order of the dims, then do a transpose
+      if (!Arrays.equals(perm, arrange)) {
+        aTrans = Transpose.create(scope, a, Constant.vectorOf(scope, perm));
+      } else {
+        aTrans = a;
+      }
+
+      // reshape the final result to the new Shape, if necessary
+      Operand<T> aReshaped =
+          aTrans.asOutput().shape().equals(newShape)
+              ? aTrans
+              : Reshape.create(scope, aTrans, Constant.vectorOf(scope, newShape.asArray()));
+      // return a tuple for the reshaped Operand, and Operand for the free dimensions, and a long
+      // array for the free dimensions
+      return new Object[] {aReshaped, Constant.vectorOf(scope, freeDims), freeDims};
+
+    } else { // calculate dynamically
+
+      long[] freeDimsStatic = null;
+      Operand<TInt32> one = Constant.scalarOf(scope, 1);
+      Operand<TInt32> minusOne = Constant.scalarOf(scope, -1);
+      Operand<TInt32> zero = Constant.scalarOf(scope, 0);
+      org.tensorflow.op.core.Shape<TInt32> tShape = org.tensorflow.op.core.Shape.create(scope, a);
+      Operand<TInt32> axesT;
+      Operand<TInt32> freeT;
+      if (aShape.numDimensions()
+          != Shape.UNKNOWN_SIZE) { // we know the rank, but there are unknown dimensions
+        long[] aShapeDims = aShape.asArray();
+        if (aShapeDims == null) aShapeDims = new long[0];
+
+        // get int array from axis Operand
+        int[] iAxes = getIntArray(scope, axis);
+        // Convert negative axes to positive
+        for (int i = 0; i < iAxes.length; i++)
+          iAxes[i] = iAxes[i] >= 0 ? iAxes[i] : Math.floorMod(iAxes[i], iAxes.length);
+
+        // convert integer axis to long axis
+        long[] lAxes = Arrays.stream(iAxes).mapToLong(i -> i).toArray();
+
+        // create list of the axes, dims, and free axes
+        List<Long> axesList = Arrays.stream(lAxes).boxed().collect(Collectors.toList());
+        List<Long> dimsList = Arrays.stream(aShapeDims).boxed().collect(Collectors.toList());
+        List<Long> freeList = new ArrayList<>(axesList);
+        freeList.removeAll(dimsList);
+
+        // create array of free dims
+        long[] freeDims = freeList.stream().mapToLong(i -> i).toArray();
+        freeDimsStatic = freeDims;
+
+        axesT = Constant.vectorOf(scope, iAxes);
+        freeT = Cast.create(scope, Constant.vectorOf(scope, freeDims), TInt32.class);
+
+      } else { // we don't know the rank yet
+        Rank rank = Rank.create(scope, a);
+
+        // convert axis to positive
+        axesT =
+            Select.create(
+                scope,
+                GreaterEqual.create(scope, axis, Constant.scalarOf(scope, 0)),
+                axis,
+                Add.create(scope, axis, rank));
+
+        SetDiff1d<TInt32, TInt32> diff =
+            SetDiff1d.create(
+                scope, Range.create(scope, Constant.scalarOf(scope, 0), rank, one), axesT);
+        freeT = diff.out();
+      }
+      Operand<TInt32> freeDims = Gather.create(scope, tShape, freeT, zero);
+      Operand<TInt32> axesDims = Gather.create(scope, tShape, axesT, zero);
+      Operand<TInt32> prodFreeDims = ReduceProd.create(scope, freeDims, minusOne);
+      Operand<TInt32> prodAxesDims = ReduceProd.create(scope, axesDims, minusOne);
+      Operand<TInt32> perm;
+      Operand<TInt32> newShape;
+      if (flipped) {
+        perm = Concat.create(scope, Arrays.asList(axesT, freeT), zero);
+        newShape = Stack.create(scope, Arrays.asList(prodAxesDims, prodFreeDims));
+      } else {
+        perm = Concat.create(scope, Arrays.asList(freeT, axesT), zero);
+        newShape = Stack.create(scope, Arrays.asList(prodFreeDims, prodAxesDims));
+      }
+      Operand<T> aReshaped = Reshape.create(scope, Transpose.create(scope, a, perm), newShape);
+      return new Object[] {aReshaped, freeDims, freeDimsStatic};
+    }
+  }
+
+  /**
+   * Gets an int array from an Operand&lt;TInt32&gt operand.
+   *
+   * @param axes the Operand to fetch the values
+   * @return the int array from an Operand&lt;TInt32&gt;
+   */
+  private static int[] getIntArray(Scope scope, Operand<TInt32> axes) {
+    List<Integer> result = new ArrayList<>();
+    if (scope.env().isEager()) {
+      axes.asTensor().scalars().forEach(s -> result.add(s.getInt()));
+    } else {
+      try (Session session = new Session((Graph) scope.env());
+          TInt32 tensor = (TInt32) session.runner().fetch(axes).run().get(0)) {
+        tensor.scalars().forEach(s -> result.add(s.getInt()));
+      }
+    }
+    return result.stream().mapToInt(i -> i).toArray();
+  }
+
+  /**
+   * Generates two sets of contraction axes for the two tensor arguments.
+   *
+   * @param a the Operand to analyze
+   * @param axis the axes
+   * @param <T> the data type for the Operand
+   * @return the contraction axes
+   */
+  @SuppressWarnings("unchecked")
+  private static <T extends TNumber> Operand<TInt32>[] tensordotAxes(
+      Scope scope, Operand<T> a, int axis) {
+    Shape aShape = a.asOutput().shape();
+    if (axis < 0) {
+      throw new IllegalArgumentException("'axis' must be at least 0.");
+    }
+    int rank = aShape.numDimensions();
+    Operand<TInt32>[] result = new Operand[2];
+    if (rank != Shape.UNKNOWN_SIZE) {
+      if (axis > rank) {
+        throw new IllegalArgumentException(
+            String.format(
+                "'axis' must not be larger than the number of dimensions of tensor %s.", rank));
+      }
+      int min = rank - axis;
+      int postRange = rank - min;
+      int[] postAxis = new int[postRange];
+      for (int i = 0; i < postRange; i++) postAxis[i] = i + min;
+
+      int[] preAxis = new int[axis];
+      for (int i = 0; i < axis; i++) preAxis[i] = i;
+
+      result[0] = Constant.vectorOf(scope, postAxis);
+      result[1] = Constant.vectorOf(scope, preAxis);
+    } else {
+      Rank rankT = Rank.create(scope, a);
+      Constant<TInt32> axisT = Constant.scalarOf(scope, axis);
+      Constant<TInt32> one = Constant.scalarOf(scope, 1);
+      Constant<TInt32> zero = Constant.scalarOf(scope, 0);
+      AssertThat assertion =
+          AssertThat.create(
+              scope,
+              Less.create(scope, axisT, rankT),
+              Arrays.asList(
+                  Constant.scalarOf(
+                      scope, "'axes' must not be larger than the number of dimensions of tensor "),
+                  rankT));
+      Scope scope1 = scope.withControlDependencies(Collections.singletonList(assertion));
+      result[0] = Range.create(scope1, Sub.create(scope, rankT, axisT), rankT, one);
+      result[1] = Range.create(scope1, zero, axisT, one);
+    }
+    return result;
+  }
+
+  /**
+   * Generates two sets of contraction axes for the two tensor arguments.
+   *
+   * @param a the Operand to analyze
+   * @param axes the axes
+   * @param <T> the data type for the Operand
+   * @return the contraction axes
+   */
+  @SuppressWarnings({"unchecked", "unused"})
+  private static <T extends TNumber> Operand<TInt32>[] tensordotAxes(
+      Scope scope, Operand<T> a, int[] axes) {
+    if (axes.length != 2)
+      throw new IllegalArgumentException(
+          "'axes' must have length 1 or 2, provided with " + axes.length);
+    int[] aAxis = new int[] {axes[0]};
+    int[] bAxis = new int[] {axes[1]};
+    Operand<TInt32>[] result = new Operand[2];
+    result[0] = Constant.vectorOf(scope, aAxis);
+    result[1] = Constant.vectorOf(scope, bAxis);
+
+    return result;
+  }
+
+  /**
+   * Generates two sets of contraction axes for the two tensor arguments.
+   *
+   * @param a the Operand to analyze
+   * @param axes the axes
+   * @param <T> the data type for the Operand
+   * @return the contraction axes
+   */
+  @SuppressWarnings({"unchecked", "unused"})
+  private static <T extends TNumber> Operand<TInt32>[] tensordotAxes(
+      Scope scope, Operand<T> a, int[][] axes) {
+    if (axes.length != 2)
+      throw new IllegalArgumentException(
+          "'axes' must have length 1 or 2, provided with " + axes.length);
+    int[] aAxis = axes[0];
+    int[] bAxis = axes[1];
+    if (aAxis.length != bAxis.length)
+      throw new IllegalArgumentException(
+          String.format(
+              "Different number of contraction axes 'a' and 'b', %d != %d",
+              aAxis.length, bAxis.length));
+    Operand<TInt32>[] result = new Operand[2];
+    result[0] = Constant.vectorOf(scope, aAxis);
+    result[1] = Constant.vectorOf(scope, bAxis);
+    return result;
+  }
+
+  /**
+   * Generates two sets of contraction axes for the two tensor arguments.
+   *
+   * @param a the Operand to analyze
+   * @param axes the axes
+   * @param <T> the data type for the Operand
+   * @return the contraction axes
+   */
+  @SuppressWarnings({"unchecked", "unused"})
+  private static <T extends TNumber> Operand<TInt32>[] tensordotAxes(
+      Scope scope, Operand<T> a, Operand<TInt32> axes) {
+
+    Constant<TInt32> one = Constant.scalarOf(scope, 1);
+    Constant<TInt32> zero = Constant.scalarOf(scope, 0);
+    Operand<TInt32>[] result = new Operand[2];
+    result[0] =
+        Slice.create(
+            scope,
+            axes,
+            Cast.create(scope, zero, TInt32.class),
+            Cast.create(scope, one, TInt32.class));
+    result[1] =
+        Slice.create(
+            scope,
+            axes,
+            Cast.create(scope, one, TInt32.class),
+            Cast.create(scope, one, TInt32.class));
+    return result;
+  }
+
+  /**
+   * Tensor contraction of a and b along specified axes and outer product.
+   *
+   * <p>Tensordot (also known as tensor contraction) sums the product of elements from {@code a} and
+   * {@code b} over the indices specified by {@code a_axes} and {@code b_axes}. The lists {@code
+   * a_axes} and {@code b_axes} specify those pairs of axes along which to contract the tensors. The
+   * axis {@code a_axes[i]} of {@code a} must have the same dimension as axis {@code b_axes[i]} of
+   * {@code b} for all {@code i} in {@code range(0, len(a_axes))}. The lists {@code a_axes} and
+   * {@code b_axes} must have identical length and consist of unique integers that specify valid
+   * axes for each of the tensors. Additionally outer product is supported by passing {@code
+   * axes=0}.
+   *
+   * <p>This operation corresponds to {@code numpy.tensordot(a, b, axes)}.
+   *
+   * <p>Example 1: When {@code a} and {@code b} are matrices (order 2), the case {@code axes = 1} is
+   * equivalent to matrix multiplication.
+   *
+   * <p>Example 2: When {@code a} and {@code b} are matrices (order 2), the case {@code axes = [[1],
+   * [0]]} is equivalent to matrix multiplication.
+   *
+   * <p>Example 3: When {@code a} and {@code b} are matrices (order 2), the case {@code axes=0}
+   * gives the outer product, a tensor of order 4.
+   *
+   * <p>Example 4: Suppose that <i>a<sub>ijk</sub></i> and <i>b<sub>lmn</sub></i> represent two
+   * tensors of order 3. Then, {@code contract(a, b, [[0], [2]])} is the order 4 tensor
+   * <i>c<sub>jklm</sub></i> whose entry corresponding to the indices <i>(j,k,l,m)</i> is given by:
+   * <i> c<sub>jklm</sub> = &Sigma;<sub>i</sub> a<sub>ijk</sub> b<sub>lmi</sub> </i>.
+   *
+   * <p>In general, {@code order(c) = order(a) + order(b) - 2*len(axes[0])}.
+   *
+   * @param a {@code Operand} of type {@code TFloat32} or {@code TFloat64}.
+   * @param b {@code Operand} with the same type as {@code a}.
+   * @param axis sum over the last N axes of a and the first N axes of b in order. If {@code
+   *     axis=0}, computes the outer product between {@code a} and {@code b}.
+   * @param <T> the datatype of the Operands, must be either TFloat32 or TFloat64
+   * @return A {@code Operand} with the same type as {@code a}.
+   * @throws IllegalArgumentException if a is not a float32 or float64 data type and if a and b are
+   *     not the same data type
+   */
+  public static <T extends TFloating> Operand<T> tensordot(
+      Scope scope, Operand<T> a, Operand<T> b, int axis) {
+
+    Operand<TInt32>[] abAxis = tensordotAxes(scope, a, axis);
+    Operand<TInt32> aAxis = abAxis[0];
+    Operand<TInt32> bAxis = abAxis[1];
+    return tensordot(scope, a, b, aAxis, bAxis);
+  }
+
+  /**
+   * Tensor contraction of a and b along specified axes and outer product.
+   *
+   * <p>Tensordot (also known as tensor contraction) sums the product of elements from {@code a} and
+   * {@code b} over the indices specified by {@code a_axes} and {@code b_axes}. The lists {@code
+   * a_axes} and {@code b_axes} specify those pairs of axes along which to contract the tensors. The
+   * axis {@code a_axes[i]} of {@code a} must have the same dimension as axis {@code b_axes[i]} of
+   * {@code b} for all {@code i} in {@code range(0, len(a_axes))}. The lists {@code a_axes} and
+   * {@code b_axes} must have identical length and consist of unique integers that specify valid
+   * axes for each of the tensors. Additionally outer product is supported by passing {@code
+   * axes=0}.
+   *
+   * <p>This operation corresponds to {@code numpy.tensordot(a, b, axes)}.
+   *
+   * <p>Example 1: When {@code a} and {@code b} are matrices (order 2), the case {@code axes = 1} is
+   * equivalent to matrix multiplication.
+   *
+   * <p>Example 2: When {@code a} and {@code b} are matrices (order 2), the case {@code axes = [[1],
+   * [0]]} is equivalent to matrix multiplication.
+   *
+   * <p>Example 3: When {@code a} and {@code b} are matrices (order 2), the case {@code axes=0}
+   * gives the outer product, a tensor of order 4.
+   *
+   * <p>Example 4: Suppose that <i>a<sub>ijk</sub></i> and <i>b<sub>lmn</sub></i> represent two
+   * tensors of order 3. Then, {@code contract(a, b, [[0], [2]])} is the order 4 tensor
+   * <i>c<sub>jklm</sub></i> whose entry corresponding to the indices <i>(j,k,l,m)</i> is given by:
+   *
+   * <p><i> c<sub>jklm</sub> = &Sigma;<sub>i</sub> a<sub>ijk</sub> b<sub>lmi</sub> </i>.
+   *
+   * <p>In general, {@code order(c) = order(a) + order(b) - 2*len(axes[0])}.
+   *
+   * <p>
+   *
+   * @param a {@code Operand} of type {@code TFloat32} or {@code TFloat64}.
+   * @param b {@code Operand} with the same type as {@code a}.
+   * @param axes If axes is a scalar, sum over the last N axes of a and the first N axes of b in
+   *     order. If axes is a list, the first and second row contain the set of unique integers
+   *     specifying axes along which the contraction is computed, for {@code a} and {@code b},
+   *     respectively. The number of axes for {@code a} and {@code b} must be equal. If {@code
+   *     axis=0}, computes the outer product between {@code a} and {@code b}.
+   * @param <T> the datatype of the Operands, must be either TFloat32 or TFloat64
+   * @return A {@code Operand} with the same type as {@code a}.
+   * @throws IllegalArgumentException if a is not a float32 or float64 data type and if a and b are
+   *     not the same data type
+   */
+  public static <T extends TFloating> Operand<T> tensordot(
+      Scope scope, Operand<T> a, Operand<T> b, Operand<TInt32> axes) {
+
+    Operand<TInt32>[] abAxis = tensordotAxes(scope, a, axes);
+    Operand<TInt32> aAxis = abAxis[0];
+    Operand<TInt32> bAxis = abAxis[1];
+
+    return tensordot(scope, a, b, aAxis, bAxis);
+  }
+
+  /**
+   * Tensor contraction of a and b along specified axes and outer product.
+   *
+   * <p>Tensordot (also known as tensor contraction) sums the product of elements from {@code a} and
+   * {@code b} over the indices specified by {@code a_axes} and {@code b_axes}. The lists {@code
+   * a_axes} and {@code b_axes} specify those pairs of axes along which to contract the tensors. The
+   * axis {@code a_axes[i]} of {@code a} must have the same dimension as axis {@code b_axes[i]} of
+   * {@code b} for all {@code i} in {@code range(0, len(a_axes))}. The lists {@code a_axes} and
+   * {@code b_axes} must have identical length and consist of unique integers that specify valid
+   * axes for each of the tensors. Additionally outer product is supported by passing {@code
+   * axes=0}.
+   *
+   * <p>This operation corresponds to {@code numpy.tensordot(a, b, axes)}.
+   *
+   * <p>Example 1: When {@code a} and {@code b} are matrices (order 2), the case {@code axes = 1} is
+   * equivalent to matrix multiplication.
+   *
+   * <p>Example 2: When {@code a} and{@code b} are matrices (order 2), the case {@code axes = [[1],
+   * [0]]} is equivalent to matrix multiplication.
+   *
+   * <p>Example 3: When {@code a} and {@code b} are matrices (order 2), the case {@code axes=0}
+   * gives the outer product, a tensor of order 4.
+   *
+   * <p>Example 4: Suppose that <i>a<sub>ijk</sub></i> and <i>b<sub>lmn</sub></i> represent two
+   * tensors of order 3. Then, {@code contract(a, b, [[0], [2]])} is the order 4 tensor
+   * <i>c<sub>jklm</sub></i> whose entry corresponding to the indices <i>(j,k,l,m)</i> is given by:
+   *
+   * <p><i> c<sub>jklm</sub> = &Sigma;<sub>i</sub> a<sub>ijk</sub> b<sub>lmi</sub> </i>.
+   *
+   * <p>In general, {@code order(c) = order(a) + order(b) - 2*len(axes[0])}.
+   *
+   * <p>
+   *
+   * @param a {@code Operand} of type {@code TFloat32} or {@code TFloat64}.
+   * @param b {@code Operand} with the same type as {@code a}.
+   * @param axes the first and second row contain the set of unique integers specifying axes along
+   *     which the contraction is computed, for {@code a} and {@code b}, respectively. The number of
+   *     axes for {@code a} and {@code b} must be equal. I
+   * @param <T> the datatype of the Operands, must be either TFloat32 or TFloat64
+   * @return A {@code Operand} with the same type as {@code a}.
+   * @throws IllegalArgumentException if a is not a float32 or float64 data type and if a and b are
+   *     not the same data type
+   */
+  public static <T extends TFloating> Operand<T> tensordot(
+      Scope scope, Operand<T> a, Operand<T> b, int[] axes) {
+
+    Operand<TInt32>[] abAxis = tensordotAxes(scope, a, axes);
+    Operand<TInt32> aAxis = abAxis[0];
+    Operand<TInt32> bAxis = abAxis[1];
+
+    return tensordot(scope, a, b, aAxis, bAxis);
+  }
+
+  /**
+   * Tensor contraction of a and b along specified axes and outer product.
+   *
+   * <p>Tensordot (also known as tensor contraction) sums the product of elements from {@code a} and
+   * {@code b} over the indices specified by {@code a_axes} and {@code b_axes}. The lists {@code
+   * a_axes} and {@code b_axes} specify those pairs of axes along which to contract the tensors. The
+   * axis {@code a_axes[i]} of {@code a} must have the same dimension as axis {@code b_axes[i]} of
+   * {@code b} for all {@code i} in {@code range(0, len(a_axes))}. The lists {@code a_axes} and
+   * {@code b_axes} must have identical length and consist of unique integers that specify valid
+   * axes for each of the tensors. Additionally outer product is supported by passing {@code
+   * axes=0}.
+   *
+   * <p>This operation corresponds to {@code numpy.tensordot(a, b, axes)}.
+   *
+   * <p>Example 1: When {@code a} and {@code b} are matrices (order 2), the case {@code axes = 1} is
+   * equivalent to matrix multiplication.
+   *
+   * <p>Example 2: When {@code a} and{@code b} are matrices (order 2), the case {@code axes = [[1],
+   * [0]]} is equivalent to matrix multiplication.
+   *
+   * <p>Example 3: When {@code a} and {@code b} are matrices (order 2), the case {@code axes=0}
+   * gives the outer product, a tensor of order 4.
+   *
+   * <p>Example 4: Suppose that <i>a<sub>ijk</sub></i> and <i>b<sub>lmn</sub></i> represent two
+   * tensors of order 3. Then, {@code contract(a, b, [[0], [2]])} is the order 4 tensor
+   * <i>c<sub>jklm</sub></i> whose entry corresponding to the indices <i>(j,k,l,m)</i> is given by:
+   *
+   * <p><i> c<sub>jklm</sub> = &Sigma;<sub>i</sub> a<sub>ijk</sub> b<sub>lmi</sub> </i>.
+   *
+   * <p>In general, {@code order(c) = order(a) + order(b) - 2*len(axes[0])}.
+   *
+   * <p>
+   *
+   * @param a {@code Operand} of type {@code TFloat32} or {@code TFloat64}.
+   * @param b {@code Operand} with the same type as {@code a}.
+   * @param axes the first and second row contain the set of unique integers specifying axes along
+   *     which the contraction is computed, for {@code a} and {@code b}, respectively. The number of
+   *     axes for {@code a} and {@code b} must be equal. I
+   * @param <T> the datatype of the Operands, must be either TFloat32 or TFloat64
+   * @return A {@code Operand} with the same type as {@code a}.
+   * @throws IllegalArgumentException if a is not a float32 or float64 data type and if a and b are
+   *     not the same data type
+   */
+  public static <T extends TFloating> Operand<T> tensordot(
+      Scope scope, Operand<T> a, Operand<T> b, int[][] axes) {
+
+    Operand<TInt32>[] abAxis = tensordotAxes(scope, a, axes);
+    Operand<TInt32> aAxis = abAxis[0];
+    Operand<TInt32> bAxis = abAxis[1];
+
+    return tensordot(scope, a, b, aAxis, bAxis);
+  }
+
+  /**
+   * Tensor contraction of a and b along specified axes and outer product.
+   *
+   * <p>Tensordot (also known as tensor contraction) sums the product of elements from {@code a} and
+   * {@code b} over the indices specified by {@code a_axes} and {@code b_axes}. The lists {@code
+   * a_axes} and {@code b_axes} specify those pairs of axes along which to contract the tensors. The
+   * axis {@code a_axes[i]} of {@code a} must have the same dimension as axis {@code b_axes[i]} of
+   * {@code b} for all {@code i} in {@code range(0, len(a_axes))}. The lists {@code a_axes} and
+   * {@code b_axes} must have identical length and consist of unique integers that specify valid
+   * axes for each of the tensors. Additionally outer product is supported by passing {@code
+   * axes=0}.
+   *
+   * <p>This operation corresponds to {@code numpy.tensordot(a, b, axes)}.
+   *
+   * <p>Example 1: When {@code a} and {@code b} are matrices (order 2), the case {@code axes = 1} is
+   * equivalent to matrix multiplication.
+   *
+   * <p>Example 2: When {@code a} and{@code b} are matrices (order 2), the case {@code axes = [[1],
+   * [0]]} is equivalent to matrix multiplication.
+   *
+   * <p>Example 3: When {@code a} and {@code b} are matrices (order 2), the case {@code axes=0}
+   * gives the outer product, a tensor of order 4.
+   *
+   * <p>Example 4: Suppose that <i>a<sub>ijk</sub></i> and <i>b<sub>lmn</sub></i> represent two
+   * tensors of order 3. Then, {@code contract(a, b, [[0], [2]])} is the order 4 tensor
+   * <i>c<sub>jklm</sub></i> whose entry corresponding to the indices <i>(j,k,l,m)</i> is given by:
+   *
+   * <p><i> c<sub>jklm</sub> = &Sigma;<sub>i</sub> a<sub>ijk</sub> b<sub>lmi</sub> </i>.
+   *
+   * <p>In general, {@code order(c) = order(a) + order(b) - 2*len(axes[0])}.
+   *
+   * <p>
+   *
+   * @param a {@code Operand} of type {@code TFloat32} or {@code TFloat64}.
+   * @param b {@code Operand} with the same type as {@code a}.
+   * @param aAxis axes for the a Operand
+   * @param bAxis axes for the b Operand
+   * @param <T> the datatype of the Operands, must be either TFloat32 or TFloat64
+   * @return A {@code Operand} with the same type as {@code a}.
+   * @throws IllegalArgumentException if a is not a float32 or float64 data type and if a and b are
+   *     not the same data type
+   */
+  @SuppressWarnings({"unchecked", "unused"})
+  public static <T extends TFloating> Operand<T> tensordot(
+      Scope scope, Operand<T> a, Operand<T> b, Operand<TInt32> aAxis, Operand<TInt32> bAxis) {
+
+    if (a.type().equals(TBfloat16.class) || a.type().equals(TFloat16.class)) {
+      throw new IllegalArgumentException(
+          String.format(
+              "Operand 'a' must be either TFloat32 or TFloat64 DataType, 'a' is a %s DataType",
+              a.type().getSimpleName()));
+    }
+    if (!a.type().equals(b.type())) {
+      throw new IllegalArgumentException(
+          String.format(
+              "Operands a and b must be the same data type, a is %s DataType, b is %s DataType",
+              a.type().getSimpleName(), b.type().getSimpleName()));
+    }
+
+    // first result is Operand<T>, second result is Operand<TIn32>, third result is long[] and it is
+    // ignored here.
+    Object[] aResult = tensordotReshape(scope, a, aAxis, false);
+    Operand<T> reshapedA = (Operand<T>) aResult[0];
+    Operand<TInt32> aFreeDims = (Operand<TInt32>) aResult[1];
+    long[] aFreeDimsStatic = (long[]) aResult[2];
+
+    // first result is Operand<T>, second result is Operand<TIn32>, third result is long[] and it is
+    // ignored here.
+    Object[] bResult = tensordotReshape(scope, b, bAxis, true);
+    Operand<T> reshapedB = (Operand<T>) bResult[0];
+    Operand<TInt32> bFreeDims = (Operand<TInt32>) bResult[1];
+    long[] bFreeDimsStatic = (long[]) bResult[2];
+
+    Operand<T> abMatmul = MatMul.matmul(scope, reshapedA, reshapedB);
+    long[] abDimsStatic = new long[aFreeDimsStatic.length + bFreeDimsStatic.length];
+    System.arraycopy(aFreeDimsStatic, 0, abDimsStatic, 0, aFreeDimsStatic.length);
+    System.arraycopy(
+        bFreeDimsStatic, 0, abDimsStatic, aFreeDimsStatic.length, bFreeDimsStatic.length);
+    if (!abMatmul.shape().hasUnknownDimension()
+        && abMatmul.shape().equals(Shape.of(abDimsStatic))) {
+      return abMatmul;
+    } else {
+      return Reshape.create(scope, abMatmul, Constant.vectorOf(scope, abDimsStatic));
+    }
+  }
+}
diff --git a/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/op/nn/SigmoidCrossEntropyWithLogits.java b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/nn/SigmoidCrossEntropyWithLogits.java
similarity index 83%
rename from tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/op/nn/SigmoidCrossEntropyWithLogits.java
rename to tensorflow-framework/src/main/java/org/tensorflow/framework/op/nn/SigmoidCrossEntropyWithLogits.java
index 92c413f7e52..b1e2ce6c928 100644
--- a/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/op/nn/SigmoidCrossEntropyWithLogits.java
+++ b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/nn/SigmoidCrossEntropyWithLogits.java
@@ -1,29 +1,31 @@
-package org.tensorflow.op.nn;
+package org.tensorflow.framework.op.nn;
 
 import org.tensorflow.Operand;
 import org.tensorflow.ndarray.Shape;
 import org.tensorflow.op.Scope;
-import org.tensorflow.op.annotation.Endpoint;
-import org.tensorflow.op.annotation.Operator;
 import org.tensorflow.op.core.Select;
 import org.tensorflow.op.core.ZerosLike;
 import org.tensorflow.op.dtypes.Cast;
-import org.tensorflow.op.math.*;
+import org.tensorflow.op.math.Add;
+import org.tensorflow.op.math.Exp;
+import org.tensorflow.op.math.GreaterEqual;
+import org.tensorflow.op.math.Log1p;
+import org.tensorflow.op.math.Mul;
+import org.tensorflow.op.math.Neg;
+import org.tensorflow.op.math.Sub;
 import org.tensorflow.types.TBool;
 import org.tensorflow.types.family.TNumber;
 
-@Operator(group = "nn")
 public class SigmoidCrossEntropyWithLogits {
 
   /**
-   * Computes sigmoid cross entropy given <code>logits</code>.
+   * Computes sigmoid cross entropy given {@code logits}.
    *
    * <p>Measures the probability error in discrete classification tasks in which each class is
    * independent and not mutually exclusive. For instance, one could perform multilabel
    * classification where a picture can contain both an elephant and a dog at the same time.
    *
-   * <p>For brevity, let <code>x = logits</code>, <code>z = labels</code>. The logistic loss in
-   * pseudo-code is
+   * <p>For brevity, let {@code x = logits}, {@code z = labels}. The logistic loss in pseudo-code is
    *
    * <pre>
    * z * -log(sigmoid(x)) + (1 - z) * -log(1 - sigmoid(x))
@@ -34,7 +36,7 @@ public class SigmoidCrossEntropyWithLogits {
    *  = x - x * z + log(1 + exp(-x))
    * </pre>
    *
-   * <p>For <code>x < 0</code>, to avoid overflow in <code>exp(-x)</code>, we reformulate the above
+   * <p>For {@code x < 0}, to avoid overflow in {@code exp(-x)}, we reformulate the above
    *
    * <pre>
    * x - x * z + log(1 + exp(-x))
@@ -49,7 +51,7 @@ public class SigmoidCrossEntropyWithLogits {
    *   max(x, 0) - x * z + log(1 + exp(-abs(x)))
    * </pre>
    *
-   * <p></ode>logits</code> and <code>labels</code> must have the same type and shape.
+   * <p>{@code logits} and {@code labels} must have the same type and shape.
    *
    * <p>
    *
@@ -60,7 +62,7 @@ public class SigmoidCrossEntropyWithLogits {
    * @return the component-wise logistic losses.
    * @throws IllegalArgumentException if logits' and labels' do not have the same shape
    */
-  @Endpoint(name = "sigmoidCrossEntropyWithLogits")
+  // @(name = "sigmoidCrossEntropyWithLogits")
   public static <T extends TNumber> Operand<T> sigmoidCrossEntropyWithLogits(
       Scope scope, Operand<T> labels, Operand<T> logits) {
     if (!isCompatible(labels.shape(), logits.shape())) {
diff --git a/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/op/nn/SoftmaxCrossEntropyWithLogits.java b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/nn/SoftmaxCrossEntropyWithLogits.java
similarity index 86%
rename from tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/op/nn/SoftmaxCrossEntropyWithLogits.java
rename to tensorflow-framework/src/main/java/org/tensorflow/framework/op/nn/SoftmaxCrossEntropyWithLogits.java
index ddeacbea4d4..a95110c9a96 100644
--- a/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/op/nn/SoftmaxCrossEntropyWithLogits.java
+++ b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/nn/SoftmaxCrossEntropyWithLogits.java
@@ -1,11 +1,17 @@
-package org.tensorflow.op.nn;
+package org.tensorflow.framework.op.nn;
 
+import java.util.Arrays;
+import java.util.List;
 import org.tensorflow.Operand;
 import org.tensorflow.ndarray.Shape;
 import org.tensorflow.op.Scope;
 import org.tensorflow.op.annotation.Endpoint;
-import org.tensorflow.op.annotation.Operator;
-import org.tensorflow.op.core.*;
+import org.tensorflow.op.core.Concat;
+import org.tensorflow.op.core.Constant;
+import org.tensorflow.op.core.Range;
+import org.tensorflow.op.core.Rank;
+import org.tensorflow.op.core.Reshape;
+import org.tensorflow.op.core.Slice;
 import org.tensorflow.op.dtypes.Cast;
 import org.tensorflow.op.linalg.Transpose;
 import org.tensorflow.op.math.Sub;
@@ -14,12 +20,7 @@
 import org.tensorflow.types.TFloat32;
 import org.tensorflow.types.TInt64;
 import org.tensorflow.types.family.TNumber;
-import org.tensorflow.types.family.TType;
-
-import java.util.Arrays;
-import java.util.List;
 
-@Operator(group = "nn")
 public class SoftmaxCrossEntropyWithLogits {
 
   /**
@@ -63,11 +64,13 @@ public class SoftmaxCrossEntropyWithLogits {
    * @param logits Per-label activations, typically a linear output. These activation energies are
    *     interpreted as unnormalized log probabilities.
    * @param axis The class dimension. -1 is the last dimension.
-   * @param <T> the number type of the operands
+   * @param <T> the data type for the logits and return operand
+   * @param <U> the data type for the labels
    * @return the softmax cross entropy loss. Its type is the same as <code>logits</code> and its
    *     shape is the same as <code>labels</code> except that it does not have the last dimension of
    *     <code>labels</code>.
    */
+  @SuppressWarnings("unchecked")
   @Endpoint(name = "softmaxCrossEntropyWithLogits")
   public static <T extends TNumber, U extends TNumber> Operand<T> softmaxCrossEntropyWithLogits(
       Scope scope, Operand<U> labels, Operand<T> logits, int axis) {
@@ -78,7 +81,9 @@ public static <T extends TNumber, U extends TNumber> Operand<T> softmaxCrossEntr
     }
 
     if (logits.asOutput().type() == TFloat16.class || logits.asOutput().type() == TBfloat16.class) {
-      Operand<TFloat32> result =  softmaxCrossEntropyWithLogits(scope,
+      Operand<TFloat32> result =
+          softmaxCrossEntropyWithLogits(
+              scope,
               Cast.create(scope, labels, TFloat32.class),
               Cast.create(scope, logits, TFloat32.class),
               axis);
@@ -86,10 +91,8 @@ public static <T extends TNumber, U extends TNumber> Operand<T> softmaxCrossEntr
     }
 
     if (logits.asOutput().type() != labels.asOutput().type()) {
-      return softmaxCrossEntropyWithLogits(scope,
-              Cast.create(scope, labels, logits.asOutput().type()),
-              logits,
-              axis);
+      return softmaxCrossEntropyWithLogits(
+          scope, Cast.create(scope, labels, logits.asOutput().type()), logits, axis);
     }
 
     Operand<TInt64> inputRank = Cast.create(scope, Rank.create(scope, logits), TInt64.class);
@@ -101,13 +104,20 @@ public static <T extends TNumber, U extends TNumber> Operand<T> softmaxCrossEntr
       labels = moveDimToEnd(scope, labels, axis, inputRank);
     }
 
+    Operand<T> tLabels;
+    if (labels.type() != logits.type()) {
+      tLabels = Cast.create(scope, labels, logits.type());
+    } else {
+      // Unchecked warning checked in if statement.
+      tLabels = (Operand<T>) labels;
+    }
+
     Shape inputShape = logits.shape();
     logits = flattenOuterDims(scope, logits);
-    labels = flattenOuterDims(scope, labels);
+    tLabels = flattenOuterDims(scope, tLabels);
 
-    org.tensorflow.op.nn.raw.SoftmaxCrossEntropyWithLogits<T> smax =
-        org.tensorflow.op.nn.raw.SoftmaxCrossEntropyWithLogits.create(
-            scope, logits, (Operand<T>)labels);
+    org.tensorflow.op.nn.SoftmaxCrossEntropyWithLogits<T> smax =
+        org.tensorflow.op.nn.SoftmaxCrossEntropyWithLogits.create(scope, logits, tLabels);
     /* cannot use generic on cost, because cost may be recast later. */
     Operand<T> cost = smax.loss();
     Operand<TInt64> outputShape =
@@ -119,6 +129,9 @@ public static <T extends TNumber, U extends TNumber> Operand<T> softmaxCrossEntr
     cost = Reshape.create(scope, cost, outputShape);
     if (scope.env().isGraph() && !shape.hasUnknownDimension()) {
       long[] array = shape.asArray();
+      if (array == null) {
+        array = new long[0];
+      }
       long[] newArray = new long[array.length - 1];
       if (axis < 0) {
         axis = shape.numDimensions() + axis;
@@ -153,7 +166,7 @@ private static <T extends TNumber> Operand<T> flattenOuterDims(Scope scope, Oper
       boolean productValid = true;
       for (int i = ndims - 2; i >= 0; i--) {
         long d = shape.size(i);
-        if (d == org.tensorflow.ndarray.Shape.UNKNOWN_SIZE) {
+        if (d == Shape.UNKNOWN_SIZE) {
           productValid = false;
           break;
         }
diff --git a/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/op/nn/SparseSoftmaxCrossEntropyWithLogits.java b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/nn/SparseSoftmaxCrossEntropyWithLogits.java
similarity index 57%
rename from tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/op/nn/SparseSoftmaxCrossEntropyWithLogits.java
rename to tensorflow-framework/src/main/java/org/tensorflow/framework/op/nn/SparseSoftmaxCrossEntropyWithLogits.java
index 54b32bb5c63..5299efcce22 100644
--- a/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/op/nn/SparseSoftmaxCrossEntropyWithLogits.java
+++ b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/nn/SparseSoftmaxCrossEntropyWithLogits.java
@@ -1,11 +1,12 @@
-package org.tensorflow.op.nn;
+package org.tensorflow.framework.op.nn;
 
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
 import org.tensorflow.Operand;
 import org.tensorflow.ndarray.Shape;
 import org.tensorflow.op.Op;
 import org.tensorflow.op.Scope;
-import org.tensorflow.op.annotation.Endpoint;
-import org.tensorflow.op.annotation.Operator;
 import org.tensorflow.op.core.AssertThat;
 import org.tensorflow.op.core.Constant;
 import org.tensorflow.op.core.Reshape;
@@ -18,15 +19,10 @@
 import org.tensorflow.types.TInt32;
 import org.tensorflow.types.family.TNumber;
 
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
-
-@Operator(group = "nn")
 public class SparseSoftmaxCrossEntropyWithLogits {
 
   /**
-   * Computes sparse softmax cross entropy between <code>logits</code> and <code>labels</code>.
+   * Computes sparse softmax cross entropy between {@code logits} and {@code labels}.
    *
    * <p>Measures the probability error in discrete classification tasks in which the classes are
    * mutually exclusive (each entry is in exactly one class). For example, each CIFAR-10 image is
@@ -35,47 +31,48 @@ public class SparseSoftmaxCrossEntropyWithLogits {
    * <p><b>NOTE:</b>
    *
    * <p>For this operation, the probability of a given label is considered exclusive. That is, soft
-   * classes are not allowed, and the <code>labels</code> vector must provide a single specific
-   * index for the true class for each row of <code>logits</code> (each minibatch entry). For soft
-   * softmax classification with a probability distribution for each entry, {@link
+   * classes are not allowed, and the {@code labels} vector must provide a single specific index for
+   * the true class for each row of {@code logits} (each minibatch entry). For soft softmax
+   * classification with a probability distribution for each entry, {@link
    * org.tensorflow.op.NnOps#softmaxCrossEntropyWithLogits}.
    *
    * <p><b>WARNING:</b>
    *
-   * <p>This op expects unscaled logits, since it performs a <code>softmax</code> on <code>logits
-   * </code> internally for efficiency. Do not call this op with the output of <code>softmax</code>,
-   * as it will produce incorrect results.
+   * <p>This op expects unscaled logits, since it performs a {@code softmax} on {@code logits }
+   * internally for efficiency. Do not call this op with the output of {@code softmax}, as it will
+   * produce incorrect results.
    *
-   * <p>A common use case is to have logits of shape <code>[batchSize, numClasses]</code> and have
-   * labels of shape <code>[batchSize]</code>, but higher dimensions are supported, in which case
-   * the <code>dim</code>-th dimension is assumed to be of size <code>numClasses</code>. <code>
-   * logits</code> must have the <cod>dataType</cod> of <code>TFloat16</code>, <code>TFloat32</code>
-   * , or <code>TFloat64</code>, and <code>labels</code> must have the dtype of <code>TInt32</code>
-   * or <code>TInt64</code>.
+   * <p>A common use case is to have logits of shape {@code [batchSize, numClasses]} and have labels
+   * of shape {@code [batchSize]}, but higher dimensions are supported, in which case the {@code
+   * dim}-th dimension is assumed to be of size {@code numClasses}. {@code logits} must have the
+   * {@code dataType} of {@code TFloat16}, {@code TFloat32} , or {@code TFloat64}, and {@code
+   * labels} must have the dtype of {@code TInt32} or {@code TInt64}.
    *
    * @param scope current scope
-   * @param labels <code>Tensor</code> of shape <code>[d_0, d_1, ..., d_{r-1}]</code> (where <code>r
-   *     </code> is rank of <code>labels</code> and result) and the dataType is <code>TInt32</code>
-   *     or <code>TInt64</code>. Each entry in <code>labels</code> must be an index in <code>[0,
-   *     numClasses)</code>. Other values will raise an exception when this op is run on CPU, and
-   *     return <code>NaN</code> for corresponding loss and gradient rows on GPU.
-   * @param logits Per-label activations (typically a linear output) of shape <code>[d_0, d_1, ...,
-   *     d_{r-1}, numClasses]</code> and dataType of <code>TFloat16</code>, <code>TFloat32</code>,
-   *     or <code>TFloat64</code>. These activation energies are interpreted as unnormalized log
-   *     probabilities.
-   * @return A <code>Tensor</code> of the same shape as <code>labels</code> and of the same type as
-   *     <code>logits</code> with the softmax cross entropy loss.
-   * @throws IllegalArgumentException If logits are scalars (need to have rank >= 1) or if the rank
-   *     of the labels is not equal to the rank of the logits minus one.
+   * @param labels {@code Tensor} of shape {@code [d_0, d_1, ..., d_{r-1}]} (where {@code r } is
+   *     rank of {@code labels} and result) and the dataType is {@code TInt32} or {@code TInt64}.
+   *     Each entry in {@code labels} must be an index in {@code [0, numClasses)}. Other values will
+   *     raise an exception when this op is run on CPU, and return {@code NaN} for corresponding
+   *     loss and gradient rows on GPU.
+   * @param logits Per-label activations (typically a linear output) of shape {@code [d_0, d_1, ...,
+   *     d_{r-1}, numClasses]} and dataType of {@code TFloat16}, {@code TFloat32}, or {@code
+   *     TFloat64}. These activation energies are interpreted as unnormalized log probabilities.
+   * @param <U> the data type for the labels
+   * @param <T> the data tyoe for the loss and logits.
+   * @return the loss
+   * @throws IllegalArgumentException If logits are scalars (need to have {@code rank >= 1}) or if
+   *     the rank of the labels is not equal to the rank of the logits minus one.
    */
-  @Endpoint(name = "sparseSoftmaxCrossEntropyWithLogits")
-  public static <T extends TNumber, U extends TNumber> Operand sparseSoftmaxCrossEntropyWithLogits(
-      Scope scope, Operand<T> labels, Operand<U> logits) {
+  @SuppressWarnings("unchecked")
+  public static <T extends TNumber, U extends TNumber>
+      Operand<T> sparseSoftmaxCrossEntropyWithLogits(
+          Scope scope, Operand<U> labels, Operand<T> logits) {
     scope = scope.withSubScope("SparseSoftmaxCrossEntropyWithLogits");
-    /** cannot use generics on preciseLogits as it may be recast later */
-    Operand preciseLogits = logits;
+    Operand<? extends TNumber> preciseLogits;
     if (logits.asOutput().type() == TFloat16.class || logits.asOutput().type() == TBfloat16.class) {
       preciseLogits = Cast.create(scope, logits, TFloat32.class);
+    } else {
+      preciseLogits = logits;
     }
     Shape labelsStaticShape = labels.shape();
     org.tensorflow.op.core.Shape<TInt32> labelsShape =
@@ -108,14 +105,16 @@ public static <T extends TNumber, U extends TNumber> Operand sparseSoftmaxCrossE
     }
     // Check if no reshapes are required.
     if (logitsShape.numDimensions() == 2) {
-      org.tensorflow.op.nn.raw.SparseSoftmaxCrossEntropyWithLogits smax =
-          org.tensorflow.op.nn.raw.SparseSoftmaxCrossEntropyWithLogits.create(
+      org.tensorflow.op.nn.SparseSoftmaxCrossEntropyWithLogits<? extends TNumber> smax =
+          org.tensorflow.op.nn.SparseSoftmaxCrossEntropyWithLogits.create(
               scope, preciseLogits, labels);
-      Operand loss = smax.loss();
-      if (logits.asOutput().type() == TFloat16.class) {
-        loss = Cast.create(scope, loss, TFloat16.class);
+      Operand<? extends TNumber> cost = smax.loss();
+      if (cost.type() != logits.type()) {
+        return Cast.create(scope, cost, logits.type());
+      } else {
+        // Unchecked cast already checked with previous if
+        return (Operand<T>) cost;
       }
-      return loss;
     }
 
     List<Op> shapeChecks = new ArrayList<>();
@@ -145,14 +144,17 @@ public static <T extends TNumber, U extends TNumber> Operand sparseSoftmaxCrossE
     preciseLogits = Reshape.create(scope, preciseLogits, Constant.arrayOf(scope, -1L, numClassses));
     labels = Reshape.create(scope, labels, Constant.scalarOf(scope, -1));
     scope.withControlDependencies(shapeChecks);
-    org.tensorflow.op.nn.raw.SparseSoftmaxCrossEntropyWithLogits smax =
-        org.tensorflow.op.nn.raw.SparseSoftmaxCrossEntropyWithLogits.create(
+    // call raw op
+    org.tensorflow.op.nn.SparseSoftmaxCrossEntropyWithLogits<? extends TNumber> smax =
+        org.tensorflow.op.nn.SparseSoftmaxCrossEntropyWithLogits.create(
             scope, preciseLogits, labels);
-    Operand cost = smax.loss();
+    Operand<? extends TNumber> cost = smax.loss();
     cost = Reshape.create(scope, cost, labelsShape);
-    if (logits.asOutput().type() == TFloat16.class) {
-      cost = Cast.create(scope, cost, TFloat16.class);
+    if (cost.type() != logits.type()) {
+      return Cast.create(scope, cost, logits.type());
+    } else {
+      // Unchecked cast already checked with previous if
+      return (Operand<T>) cost;
     }
-    return cost;
   }
 }
diff --git a/tensorflow-framework/src/main/java/org/tensorflow/framework/op/sets/Sets.java b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/sets/Sets.java
new file mode 100644
index 00000000000..0cbfa74e770
--- /dev/null
+++ b/tensorflow-framework/src/main/java/org/tensorflow/framework/op/sets/Sets.java
@@ -0,0 +1,148 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+=======================================================================*/
+package org.tensorflow.framework.op.sets;
+
+import org.tensorflow.Operand;
+import org.tensorflow.op.Scope;
+import org.tensorflow.op.SparseOps;
+import org.tensorflow.op.core.Constant;
+import org.tensorflow.op.dtypes.Cast;
+import org.tensorflow.op.sparse.DenseToDenseSetOperation;
+import org.tensorflow.op.sparse.SparseToDense;
+import org.tensorflow.types.family.TNumber;
+
+public class Sets {
+
+  /**
+   * Computes set difference of elements in last dimension of <code>a</code> and <code>b</code> with
+   * <code>aMinusB</code> set to true.
+   *
+   * <p>All but the last dimension of <code>a</code> and <code>b</code> must match
+   *
+   * @param a The first operand representing set <code>a</code>
+   * @param b The other operand representing set <code>b</code>
+   * @param <T>the data type for the sets
+   * @return An Operand with the same rank as <code>a</code> and <code>b</code>, and all but the
+   *     last dimension the * same. Elements along the last dimension contain the results of the set
+   *     operation.
+   */
+  public static <T extends TNumber> Operand<T> difference(Scope scope, Operand<T> a, Operand<T> b) {
+    return difference(scope, a, b, true);
+  }
+
+  /**
+   * Computes set difference of elements in last dimension of <code>a</code> and <code>b</code>.
+   *
+   * <p>All but the last dimension of <code>a</code> and <code>b</code> must match
+   *
+   * @param a The first operand representing set <code>a</code>
+   * @param b The other operand representing set <code>b</code>
+   * @param aMinusB whether to subtract b from a, vs vice versa.
+   * @param <T>the data type for the sets
+   * @return An Operand with the same rank as <code>a</code> and <code>b</code>, and all but the
+   *     last dimension the * same. Elements along the last dimension contain the results of the set
+   *     operation.
+   */
+  public static <T extends TNumber> Operand<T> difference(
+      Scope scope, Operand<T> a, Operand<T> b, boolean aMinusB) {
+    return setOperation(scope, a, b, aMinusB ? Operation.A_MINUS_B : Operation.B_MINUS_A);
+  }
+
+  /**
+   * Computes set union of elements in last dimension of <code>a</code> and <code>b</code>.
+   *
+   * @param a The first operand representing set <code>a</code>
+   * @param b The other operand representing set <code>b</code>
+   * @param <T>the data type for the sets
+   * @return An Operand with the same rank as <code>a</code> and <code>b</code>, and all but the
+   *     last dimension the * same. Elements along the last dimension contain the results of the set
+   *     operation.
+   */
+  public static <T extends TNumber> Operand<T> union(Scope scope, Operand<T> a, Operand<T> b) {
+    return setOperation(scope, a, b, Operation.UNION);
+  }
+
+  /**
+   * Computes set intersection of elements in last dimension of <code>a</code> and <code>b</code>.
+   *
+   * @param a The first operand representing set <code>a</code>
+   * @param b The other operand representing set <code>b</code>
+   * @param <T>the data type for the sets
+   * @return An Operand with the same rank as <code>a</code> and <code>b</code>, and all but the
+   *     last dimension the * same. Elements along the last dimension contain the results of the set
+   *     operation.
+   */
+  public static <T extends TNumber> Operand<T> intersection(
+      Scope scope, Operand<T> a, Operand<T> b) {
+    return setOperation(scope, a, b, Operation.INTERSECTION);
+  }
+
+  /**
+   * Compute set operation of elements in last dimension of <code>a</code> and <code>b</code>.
+   *
+   * @param a The first set operation operand
+   * @param b The other et operation operand
+   * @param setOperation The set operation to perform, {@link Operation}.
+   * @param <T> the data type for the sets
+   * @return An Operand with the same rank as <code>a</code> and <code>b</code>, and all but the
+   *     last dimension the same. Elements along the last dimension contain the results of the set
+   *     operation.
+   */
+  public static <T extends TNumber> Operand<T> setOperation(
+      Scope scope, Operand<T> a, Operand<T> b, Operation setOperation) {
+
+    DenseToDenseSetOperation<T> setOperationResult =
+        DenseToDenseSetOperation.create(
+            scope,
+            a,
+            b,
+            setOperation.getSetOperation(),
+            DenseToDenseSetOperation.validateIndices(true));
+
+    return SparseToDense.create(
+        scope,
+        setOperationResult.resultIndices(),
+        setOperationResult.resultShape(),
+        setOperationResult.resultValues(),
+        Cast.create(scope, Constant.scalarOf(scope, 0), a.type()));
+  }
+
+  /**
+   * Enumeration containing the string operation values to be passed to the TensorFlow Sparse Ops
+   * function {@link SparseOps#denseToDenseSetOperation}
+   */
+  public enum Operation {
+    A_MINUS_B("a-b"),
+    B_MINUS_A("b-a"),
+    INTERSECTION("intersection"),
+    UNION("union");
+
+    private final String setOperation;
+
+    Operation(String setOperation) {
+      this.setOperation = setOperation;
+    }
+
+    /**
+     * Gets the set operation String value used to pass as the stringOperation value to {@link
+     * SparseOps#denseToDenseSetOperation}
+     *
+     * @return the set operation String value
+     */
+    public String getSetOperation() {
+      return setOperation;
+    }
+  }
+}
diff --git a/tensorflow-framework/src/test/java/org/tensorflow/framework/op/LinalgOpsTest.java b/tensorflow-framework/src/test/java/org/tensorflow/framework/op/LinalgOpsTest.java
new file mode 100644
index 00000000000..f2c297ce032
--- /dev/null
+++ b/tensorflow-framework/src/test/java/org/tensorflow/framework/op/LinalgOpsTest.java
@@ -0,0 +1,60 @@
+package org.tensorflow.framework.op;
+
+import org.junit.jupiter.api.Test;
+import org.tensorflow.Operand;
+import org.tensorflow.framework.utils.TestSession;
+import org.tensorflow.op.Ops;
+import org.tensorflow.types.TFloat32;
+import org.tensorflow.types.TFloat64;
+
+class LinalgOpsTest {
+  private final TestSession.Mode[] tfModes = {TestSession.Mode.EAGER, TestSession.Mode.GRAPH};
+
+  @Test
+  public void test2D() {
+    for (TestSession.Mode tfMode : tfModes)
+      try (TestSession session = TestSession.createTestSession(tfMode)) {
+        Ops tf = session.getTF();
+        FrameworkOps fops = FrameworkOps.create(tf);
+        Operand<TFloat32> a = tf.constant(new float[][] {{3.7213619f}});
+        Operand<TFloat32> b = tf.constant(new float[][] {{8.153921f}});
+
+        Operand<TFloat32> ans = fops.linalg.matmul(a, b);
+        Operand<TFloat32> expected = tf.constant(new float[][] {{30.34369f}});
+        session.evaluate(expected, ans);
+
+        Operand<TFloat64> a64 =
+            tf.constant(new double[][] {{-8.944851}, {4.1711287}, {-0.22380222}});
+        Operand<TFloat64> b64 =
+            tf.constant(
+                new double[][] {{-14.276086, -12.433481, -2.2447076, -1.5775859, 1.8588694}});
+
+        Operand<TFloat64> ans64 = fops.linalg.matmul(a64, b64);
+        Operand<TFloat64> expected64 =
+            tf.constant(
+                new double[][] {
+                  {127.69746, 111.21564, 20.078575, 14.111271, -16.62731},
+                  {-59.547394, -51.861652, -9.362965, -6.580314, 7.753584},
+                  {3.1950197, 2.7826407, 0.50237054, 0.35306725, -0.4160191}
+                });
+        session.evaluate(expected64, ans64);
+
+        a64 =
+            tf.constant(
+                new double[][] {
+                  {-9.189821, -1.588742, -8.684379},
+                  {-10.953391, -8.473055, -6.8909864},
+                  {-11.712155, -6.6350083, -2.4441578},
+                  {1.4037079, -11.279383, 0.9129576},
+                  {0.11368857, 2.3792067, -11.218701},
+                });
+        b64 = tf.constant(new double[][] {{-4.933953}, {-12.692161}, {-10.192119}});
+        ans64 = fops.linalg.matmul(a64, b64);
+        expected64 =
+            tf.constant(
+                new double[][] {{154.01892}, {231.81863}, {166.91096}, {126.92895}, {83.58413}});
+        session.setEpsilon(1e-4f);
+        session.evaluate(expected64, ans64);
+      }
+  }
+}
diff --git a/tensorflow-framework/src/test/java/org/tensorflow/framework/op/MathOpsTest.java b/tensorflow-framework/src/test/java/org/tensorflow/framework/op/MathOpsTest.java
new file mode 100644
index 00000000000..852dc0f5dcc
--- /dev/null
+++ b/tensorflow-framework/src/test/java/org/tensorflow/framework/op/MathOpsTest.java
@@ -0,0 +1,499 @@
+package org.tensorflow.framework.op;
+
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import org.junit.jupiter.api.Test;
+import org.tensorflow.Operand;
+import org.tensorflow.framework.utils.TestSession;
+import org.tensorflow.ndarray.Shape;
+import org.tensorflow.op.Ops;
+import org.tensorflow.types.TFloat32;
+import org.tensorflow.types.TFloat64;
+import org.tensorflow.types.TInt64;
+
+class MathOpsTest {
+
+  private final TestSession.Mode[] tfModes = {TestSession.Mode.EAGER, TestSession.Mode.GRAPH};
+
+  double[][][] array =
+      new double[][][] {
+        {
+          {4.17021990e-01, 7.20324516e-01, 1.14374816e-04},
+          {3.02332580e-01, 1.46755889e-01, 9.23385918e-02},
+          {1.86260208e-01, 3.45560730e-01, 3.96767467e-01},
+          {5.38816750e-01, 4.19194520e-01, 6.85219526e-01},
+          {2.04452246e-01, 8.78117442e-01, 2.73875929e-02},
+          {6.70467496e-01, 4.17304814e-01, 5.58689833e-01},
+          {1.40386939e-01, 1.98101491e-01, 8.00744593e-01}
+        },
+        {
+          {9.68261600e-01, 3.13424170e-01, 6.92322612e-01},
+          {8.76389146e-01, 8.94606650e-01, 8.50442126e-02},
+          {3.90547849e-02, 1.69830427e-01, 8.78142476e-01},
+          {9.83468369e-02, 4.21107620e-01, 9.57889557e-01},
+          {5.33165276e-01, 6.91877127e-01, 3.15515637e-01},
+          {6.86500907e-01, 8.34625661e-01, 1.82882771e-02},
+          {7.50144303e-01, 9.88861084e-01, 7.48165667e-01}
+        },
+        {
+          {2.80443996e-01, 7.89279342e-01, 1.03226006e-01},
+          {4.47893530e-01, 9.08595502e-01, 2.93614149e-01},
+          {2.87775338e-01, 1.30028576e-01, 1.93669572e-02},
+          {6.78835511e-01, 2.11628109e-01, 2.65546650e-01},
+          {4.91573155e-01, 5.33625446e-02, 5.74117601e-01},
+          {1.46728575e-01, 5.89305520e-01, 6.99758351e-01},
+          {1.02334432e-01, 4.14055973e-01, 6.94400132e-01}
+        },
+        {
+          {4.14179265e-01, 4.99534607e-02, 5.35896420e-01},
+          {6.63794637e-01, 5.14889121e-01, 9.44594741e-01},
+          {5.86555064e-01, 9.03401911e-01, 1.37474701e-01},
+          {1.39276341e-01, 8.07391286e-01, 3.97676826e-01},
+          {1.65354192e-01, 9.27508593e-01, 3.47765863e-01},
+          {7.50812113e-01, 7.25997984e-01, 8.83306086e-01},
+          {6.23672187e-01, 7.50942409e-01, 3.48898351e-01}
+        },
+        {
+          {2.69927889e-01, 8.95886242e-01, 4.28091198e-01},
+          {9.64840055e-01, 6.63441479e-01, 6.21695697e-01},
+          {1.14745975e-01, 9.49489236e-01, 4.49912131e-01},
+          {5.78389585e-01, 4.08136815e-01, 2.37026975e-01},
+          {9.03379500e-01, 5.73679507e-01, 2.87032709e-03},
+          {6.17144942e-01, 3.26644897e-01, 5.27058125e-01},
+          {8.85942101e-01, 3.57269764e-01, 9.08535123e-01}
+        },
+        {
+          {6.23360097e-01, 1.58212427e-02, 9.29437220e-01},
+          {6.90896928e-01, 9.97322857e-01, 1.72340512e-01},
+          {1.37135744e-01, 9.32595491e-01, 6.96818173e-01},
+          {6.60001710e-02, 7.55463064e-01, 7.53876209e-01},
+          {9.23024535e-01, 7.11524785e-01, 1.24270961e-01},
+          {1.98801346e-02, 2.62109861e-02, 2.83064879e-02},
+          {2.46211067e-01, 8.60027969e-01, 5.38831055e-01}
+        },
+        {
+          {5.52821994e-01, 8.42030883e-01, 1.24173313e-01},
+          {2.79183686e-01, 5.85759282e-01, 9.69595730e-01},
+          {5.61030209e-01, 1.86472889e-02, 8.00632656e-01},
+          {2.32974276e-01, 8.07105184e-01, 3.87860656e-01},
+          {8.63541842e-01, 7.47121632e-01, 5.56240261e-01},
+          {1.36455223e-01, 5.99176884e-02, 1.21343456e-01},
+          {4.45518792e-02, 1.07494131e-01, 2.25709334e-01}
+        },
+        {
+          {7.12988973e-01, 5.59717000e-01, 1.25559801e-02},
+          {7.19742775e-02, 9.67276335e-01, 5.68100452e-01},
+          {2.03293234e-01, 2.52325743e-01, 7.43825853e-01},
+          {1.95429474e-01, 5.81358910e-01, 9.70019996e-01},
+          {8.46828818e-01, 2.39847764e-01, 4.93769705e-01},
+          {6.19955719e-01, 8.28980923e-01, 1.56791389e-01},
+          {1.85762029e-02, 7.00221434e-02, 4.86345112e-01}
+        },
+        {
+          {6.06329441e-01, 5.68851411e-01, 3.17362398e-01},
+          {9.88616168e-01, 5.79745233e-01, 3.80141169e-01},
+          {5.50948203e-01, 7.45334446e-01, 6.69232905e-01},
+          {2.64919549e-01, 6.63348362e-02, 3.70084196e-01},
+          {6.29717529e-01, 2.10174009e-01, 7.52755582e-01},
+          {6.65364787e-02, 2.60315090e-01, 8.04754555e-01},
+          {1.93434283e-01, 6.39460862e-01, 5.24670303e-01}
+        },
+        {
+          {9.24807966e-01, 2.63296783e-01, 6.59610927e-02},
+          {7.35065937e-01, 7.72178054e-01, 9.07815874e-01},
+          {9.31972086e-01, 1.39515726e-02, 2.34362081e-01},
+          {6.16778374e-01, 9.49016333e-01, 9.50176120e-01},
+          {5.56653202e-01, 9.15606380e-01, 6.41566217e-01},
+          {3.90007704e-01, 4.85990673e-01, 6.04310513e-01},
+          {5.49547911e-01, 9.26181436e-01, 9.18733418e-01}
+        },
+        {
+          {3.94875616e-01, 9.63262558e-01, 1.73955664e-01},
+          {1.26329526e-01, 1.35079160e-01, 5.05662143e-01},
+          {2.15248056e-02, 9.47970212e-01, 8.27115476e-01},
+          {1.50189810e-02, 1.76196262e-01, 3.32063586e-01},
+          {1.30996838e-01, 8.09490681e-01, 3.44736665e-01},
+          {9.40107465e-01, 5.82014203e-01, 8.78831983e-01},
+          {8.44734430e-01, 9.05392289e-01, 4.59880263e-01}
+        },
+        {
+          {5.46346843e-01, 7.98603594e-01, 2.85718858e-01},
+          {4.90253508e-01, 5.99110305e-01, 1.55332759e-02},
+          {5.93481421e-01, 4.33676362e-01, 8.07360530e-01},
+          {3.15244794e-01, 8.92888725e-01, 5.77857196e-01},
+          {1.84010208e-01, 7.87929237e-01, 6.12031162e-01},
+          {5.39092720e-02, 4.20193672e-01, 6.79068863e-01},
+          {9.18601751e-01, 4.02024889e-04, 9.76759136e-01}
+        },
+        {
+          {3.76580328e-01, 9.73783553e-01, 6.04716122e-01},
+          {8.28845799e-01, 5.74711502e-01, 6.28076196e-01},
+          {2.85576284e-01, 5.86833358e-01, 7.50021756e-01},
+          {8.58313859e-01, 7.55082190e-01, 6.98057234e-01},
+          {8.64479423e-01, 3.22681010e-01, 6.70788765e-01},
+          {4.50873941e-01, 3.82102758e-01, 4.10811365e-01},
+          {4.01479572e-01, 3.17383945e-01, 6.21919394e-01}
+        },
+        {
+          {4.30247277e-01, 9.73802090e-01, 6.77800894e-01},
+          {1.98569894e-01, 4.26701009e-01, 3.43346238e-01},
+          {7.97638834e-01, 8.79998267e-01, 9.03841972e-01},
+          {6.62719786e-01, 2.70208269e-01, 2.52366692e-01},
+          {8.54897916e-01, 5.27714670e-01, 8.02161098e-01},
+          {5.72488546e-01, 7.33142555e-01, 5.19011617e-01},
+          {7.70883918e-01, 5.68857968e-01, 4.65709865e-01}
+        },
+        {
+          {3.42688918e-01, 6.82093501e-02, 3.77924174e-01},
+          {7.96260759e-02, 9.82817113e-01, 1.81612849e-01},
+          {8.11858714e-01, 8.74961674e-01, 6.88413262e-01},
+          {5.69494426e-01, 1.60971433e-01, 4.66880023e-01},
+          {3.45172048e-01, 2.25039959e-01, 5.92511892e-01},
+          {3.12269837e-01, 9.16305542e-01, 9.09635544e-01},
+          {2.57118285e-01, 1.10891297e-01, 1.92962736e-01}
+        },
+        {
+          {4.99584168e-01, 7.28585660e-01, 2.08194435e-01},
+          {2.48033553e-01, 8.51671875e-01, 4.15848732e-01},
+          {6.16685092e-01, 2.33666137e-01, 1.01967260e-01},
+          {5.15857041e-01, 4.77140993e-01, 1.52671650e-01},
+          {6.21806204e-01, 5.44010103e-01, 6.54137373e-01},
+          {1.44545540e-01, 7.51527846e-01, 2.22049147e-01},
+          {5.19351840e-01, 7.85296023e-01, 2.23304275e-02}
+        },
+        {
+          {3.24362457e-01, 8.72922361e-01, 8.44709635e-01},
+          {5.38440585e-01, 8.66608262e-01, 9.49805975e-01},
+          {8.26407015e-01, 8.54115427e-01, 9.87434015e-02},
+          {6.51304305e-01, 7.03516960e-01, 6.10240817e-01},
+          {7.99615264e-01, 3.45712192e-02, 7.70238757e-01},
+          {7.31728613e-01, 2.59698391e-01, 2.57069290e-01},
+          {6.32303298e-01, 3.45297456e-01, 7.96588659e-01}
+        },
+        {
+          {4.46146220e-01, 7.82749414e-01, 9.90471780e-01},
+          {3.00248325e-01, 1.43005833e-01, 9.01308417e-01},
+          {5.41559398e-01, 9.74740386e-01, 6.36604428e-01},
+          {9.93912995e-01, 5.46070814e-01, 5.26425958e-01},
+          {1.35427907e-01, 3.55705172e-01, 2.62185670e-02},
+          {1.60395175e-01, 7.45637178e-01, 3.03996895e-02},
+          {3.66543084e-01, 8.62346232e-01, 6.92677736e-01}
+        },
+        {
+          {6.90942168e-01, 1.88636795e-01, 4.41904277e-01},
+          {5.81577420e-01, 9.89751697e-01, 2.03906223e-01},
+          {2.47732908e-01, 2.62173086e-01, 7.50172436e-01},
+          {4.56975341e-01, 5.69294393e-02, 5.08516252e-01},
+          {2.11960167e-01, 7.98604250e-01, 2.97331393e-01},
+          {2.76060123e-02, 5.93432426e-01, 8.43840420e-01},
+          {3.81016135e-01, 7.49858320e-01, 5.11141479e-01}
+        },
+        {
+          {5.40951788e-01, 9.59434330e-01, 8.03960919e-01},
+          {3.23230661e-02, 7.09387243e-01, 4.65001494e-01},
+          {9.47548926e-01, 2.21432731e-01, 2.67072022e-01},
+          {8.14739615e-02, 4.28618819e-01, 1.09018765e-01},
+          {6.33786738e-01, 8.02963257e-01, 6.96800470e-01},
+          {7.66211390e-01, 3.42454106e-01, 8.45851481e-01},
+          {4.28768784e-01, 8.24009895e-01, 6.26496136e-01}
+        }
+      };
+
+  double[][][] expectedArray = {
+    {
+      {3.45350616e-02, 5.96526116e-02, 9.47178160e-06},
+      {2.50372272e-02, 1.21533722e-02, 7.64688430e-03},
+      {1.54248644e-02, 2.86171008e-02, 3.28577124e-02},
+      {4.46213149e-02, 3.47149745e-02, 5.67454435e-02},
+      {1.69314109e-02, 7.27199987e-02, 2.26806314e-03},
+      {5.55237755e-02, 3.45584825e-02, 4.62670736e-02},
+      {1.16259372e-02, 1.64054818e-02, 6.63124844e-02}
+    },
+    {
+      {8.01851526e-02, 2.59557609e-02, 5.73336743e-02},
+      {7.25768730e-02, 7.40855262e-02, 7.04281079e-03},
+      {3.23426444e-03, 1.40642561e-02, 7.27220699e-02},
+      {8.14444851e-03, 3.48734073e-02, 7.93262124e-02},
+      {4.41532955e-02, 5.72967827e-02, 2.61289626e-02},
+      {5.68515584e-02, 6.91182911e-02, 1.51451665e-03},
+      {6.21220917e-02, 8.18910673e-02, 6.19582348e-02}
+    },
+    {
+      {2.32245550e-02, 6.53630048e-02, 8.54850933e-03},
+      {3.70916426e-02, 7.52439946e-02, 2.43152231e-02},
+      {2.38316897e-02, 1.07681248e-02, 1.60384597e-03},
+      {5.62167615e-02, 1.75256692e-02, 2.19908543e-02},
+      {4.07089069e-02, 4.41914052e-03, 4.75447029e-02},
+      {1.21511100e-02, 4.88024652e-02, 5.79494536e-02},
+      {8.47467501e-03, 3.42894346e-02, 5.75057231e-02}
+    },
+    {
+      {3.42996456e-02, 4.13682219e-03, 4.43794727e-02},
+      {5.49711734e-02, 4.26397808e-02, 7.82252178e-02},
+      {4.85746935e-02, 7.48138949e-02, 1.13847647e-02},
+      {1.15339644e-02, 6.68629184e-02, 3.29330191e-02},
+      {1.36935636e-02, 7.68102556e-02, 2.87997164e-02},
+      {6.21773973e-02, 6.01224527e-02, 7.31496885e-02},
+      {5.16484901e-02, 6.21881858e-02, 2.88935024e-02}
+    },
+    {
+      {2.23536789e-02, 7.41914958e-02, 3.54517400e-02},
+      {7.99018070e-02, 5.49419262e-02, 5.14848121e-02},
+      {9.50251892e-03, 7.86305517e-02, 3.72588076e-02},
+      {4.78984788e-02, 3.37992460e-02, 1.96290389e-02},
+      {7.48120397e-02, 4.75084223e-02, 2.37701897e-04},
+      {5.11079468e-02, 2.70506144e-02, 4.36475389e-02},
+      {7.33679906e-02, 2.95867678e-02, 7.52389953e-02}
+    },
+    {
+      {5.16226478e-02, 1.31021289e-03, 7.69699737e-02},
+      {5.72156087e-02, 8.25918168e-02, 1.42721254e-02},
+      {1.13566946e-02, 7.72315189e-02, 5.77059686e-02},
+      {5.46570681e-03, 6.25625551e-02, 6.24311455e-02},
+      {7.64389113e-02, 5.89238741e-02, 1.02913165e-02},
+      {1.64634397e-03, 2.17062421e-03, 2.34416011e-03},
+      {2.03896053e-02, 7.12219477e-02, 4.46224995e-02}
+    },
+    {
+      {4.57811356e-02, 6.97315410e-02, 1.02832299e-02},
+      {2.31201854e-02, 4.85087894e-02, 8.02956372e-02},
+      {4.64608893e-02, 1.54424773e-03, 6.63032085e-02},
+      {1.92934200e-02, 6.68392256e-02, 3.21201086e-02},
+      {7.15129450e-02, 6.18717745e-02, 4.60642166e-02},
+      {1.13003375e-02, 4.96199494e-03, 1.00488793e-02},
+      {3.68949817e-03, 8.90196767e-03, 1.86917856e-02}
+    },
+    {
+      {5.90451285e-02, 4.63521369e-02, 1.03980501e-03},
+      {5.96044352e-03, 8.01035613e-02, 4.70464006e-02},
+      {1.68354288e-02, 2.08959840e-02, 6.15988411e-02},
+      {1.61842033e-02, 4.81443815e-02, 8.03307742e-02},
+      {7.01288804e-02, 1.98626388e-02, 4.08908091e-02},
+      {5.13407178e-02, 6.86508343e-02, 1.29844472e-02},
+      {1.53836084e-03, 5.79878036e-03, 4.02759537e-02}
+    },
+    {
+      {5.02122790e-02, 4.71085906e-02, 2.62818988e-02},
+      {8.18707868e-02, 4.80107442e-02, 3.14808302e-02},
+      {4.56259623e-02, 6.17237724e-02, 5.54215349e-02},
+      {2.19389219e-02, 5.49342157e-03, 3.06479763e-02},
+      {5.21491282e-02, 1.74052510e-02, 6.23383410e-02},
+      {5.51012019e-03, 2.15576105e-02, 6.66445568e-02},
+      {1.60189737e-02, 5.29560074e-02, 4.34497967e-02}
+    },
+    {
+      {7.65866041e-02, 2.18045339e-02, 5.46247046e-03},
+      {6.08734004e-02, 6.39467835e-02, 7.51794279e-02},
+      {7.71798939e-02, 1.15537888e-03, 1.94083489e-02},
+      {5.10775894e-02, 7.85913840e-02, 7.86874294e-02},
+      {4.60984148e-02, 7.58245885e-02, 5.31303585e-02},
+      {3.22979130e-02, 4.02465984e-02, 5.00450842e-02},
+      {4.55099978e-02, 7.67003447e-02, 7.60835484e-02}
+    },
+    {
+      {3.27010415e-02, 7.97711685e-02, 1.44058811e-02},
+      {1.04617933e-02, 1.11863809e-02, 4.18756641e-02},
+      {1.78254500e-03, 7.85047561e-02, 6.84963465e-02},
+      {1.24377478e-03, 1.45914331e-02, 2.74993554e-02},
+      {1.08483098e-02, 6.70367777e-02, 2.85488572e-02},
+      {7.78536126e-02, 4.81986478e-02, 7.27791712e-02},
+      {6.99554384e-02, 7.49787241e-02, 3.80843058e-02}
+    },
+    {
+      {4.52449061e-02, 6.61351755e-02, 2.36613862e-02},
+      {4.05996218e-02, 4.96144369e-02, 1.28636532e-03},
+      {4.91482876e-02, 3.59142683e-02, 6.68603703e-02},
+      {2.61065327e-02, 7.39432648e-02, 4.78543900e-02},
+      {1.52385337e-02, 6.52511939e-02, 5.06844558e-02},
+      {4.46441676e-03, 3.47977169e-02, 5.62360846e-02},
+      {7.60726482e-02, 3.32930977e-05, 8.08888674e-02}
+    },
+    {
+      {3.11859436e-02, 8.06424469e-02, 5.00786714e-02},
+      {6.86396435e-02, 4.75938842e-02, 5.20132035e-02},
+      {2.36495789e-02, 4.85977381e-02, 6.21119440e-02},
+      {7.10799918e-02, 6.25310168e-02, 5.78085780e-02},
+      {7.15905875e-02, 2.67223511e-02, 5.55503815e-02},
+      {3.73384580e-02, 3.16432752e-02, 3.40207368e-02},
+      {3.32479365e-02, 2.62836833e-02, 5.15033379e-02}
+    },
+    {
+      {3.56302932e-02, 8.06439817e-02, 5.61310798e-02},
+      {1.64442733e-02, 3.53366137e-02, 2.84337122e-02},
+      {6.60552830e-02, 7.28757605e-02, 7.48503357e-02},
+      {5.48821613e-02, 2.23768987e-02, 2.08993759e-02},
+      {7.07971081e-02, 4.37019095e-02, 6.64297864e-02},
+      {4.74097952e-02, 6.07141182e-02, 4.29811813e-02},
+      {6.38396144e-02, 4.71091345e-02, 3.85670736e-02}
+    },
+    {
+      {2.83792764e-02, 5.64865675e-03, 3.12972330e-02},
+      {6.59411587e-03, 8.13905448e-02, 1.50400000e-02},
+      {6.72328845e-02, 7.24586621e-02, 5.70099279e-02},
+      {4.71618399e-02, 1.33306114e-02, 3.86639796e-02},
+      {2.85849143e-02, 1.86363515e-02, 4.90679964e-02},
+      {2.58601662e-02, 7.58824944e-02, 7.53301233e-02},
+      {2.12928709e-02, 9.18329880e-03, 1.59799233e-02}
+    },
+    {
+      {4.13723253e-02, 6.03367463e-02, 1.72413141e-02},
+      {2.05405317e-02, 7.05299526e-02, 3.44378985e-02},
+      {5.10698669e-02, 1.93507168e-02, 8.44426826e-03},
+      {4.27199379e-02, 3.95137258e-02, 1.26432776e-02},
+      {5.14939614e-02, 4.50513922e-02, 5.41714206e-02},
+      {1.19703254e-02, 6.22366704e-02, 1.83886718e-02},
+      {4.30093557e-02, 6.50331303e-02, 1.84926135e-03}
+    },
+    {
+      {2.68615987e-02, 7.22897798e-02, 6.99533820e-02},
+      {4.45901640e-02, 7.17668831e-02, 7.86567777e-02},
+      {6.84376806e-02, 7.07323104e-02, 8.17728881e-03},
+      {5.39368056e-02, 5.82607202e-02, 5.05361930e-02},
+      {6.62189573e-02, 2.86296452e-03, 6.37861863e-02},
+      {6.05970249e-02, 2.15065386e-02, 2.12888140e-02},
+      {5.23632653e-02, 2.85952985e-02, 6.59683123e-02}
+    },
+    {
+      {3.69469412e-02, 6.48222342e-02, 8.20244551e-02},
+      {2.48646215e-02, 1.18428171e-02, 7.46405274e-02},
+      {4.48484421e-02, 8.07216838e-02, 5.27194552e-02},
+      {8.23094398e-02, 4.52220477e-02, 4.35951874e-02},
+      {1.12152621e-02, 2.94571985e-02, 2.17125192e-03},
+      {1.32828895e-02, 6.17488436e-02, 2.51750532e-03},
+      {3.03547252e-02, 7.14139268e-02, 5.73630854e-02}
+    },
+    {
+      {5.72193563e-02, 1.56216780e-02, 3.65956500e-02},
+      {4.81624752e-02, 8.19648281e-02, 1.68861933e-02},
+      {2.05156356e-02, 2.17114780e-02, 6.21244237e-02},
+      {3.78437378e-02, 4.71452763e-03, 4.21120226e-02},
+      {1.75531674e-02, 6.61352351e-02, 2.46230606e-02},
+      {2.28615105e-03, 4.91442308e-02, 6.98814020e-02},
+      {3.15532871e-02, 6.20984100e-02, 4.23294269e-02}
+    },
+    {
+      {4.47981246e-02, 7.94541389e-02, 6.65788352e-02},
+      {2.67678709e-03, 5.87468557e-02, 3.85084115e-02},
+      {7.84698650e-02, 1.83376241e-02, 2.21171752e-02},
+      {6.74714567e-03, 3.54954340e-02, 9.02822800e-03},
+      {5.24861142e-02, 6.64962158e-02, 5.77045009e-02},
+      {6.34526685e-02, 2.83598304e-02, 7.00479448e-02},
+      {3.55078541e-02, 6.82391599e-02, 5.18823527e-02}
+    }
+  };
+
+  @Test
+  public void testL2Normalize() {
+    for (TestSession.Mode tfMode : tfModes)
+      try (TestSession session = TestSession.createTestSession(tfMode)) {
+        Ops tf = session.getTF();
+        FrameworkOps fops = FrameworkOps.create(tf);
+        Operand<TFloat64> input = tf.constant(array);
+        Operand<TFloat64> result = fops.math.l2Normalize(tf.constant(array), new int[] {0, 1, 2});
+        session.evaluate(tf.constant(expectedArray), result);
+      }
+  }
+
+  @Test
+  public void testConfusionMatrix() {
+    for (TestSession.Mode tfMode : tfModes)
+      try (TestSession session = TestSession.createTestSession(tfMode)) {
+        Ops tf = session.getTF();
+        FrameworkOps fops = FrameworkOps.create(tf);
+        long[] labels = new long[] {2, 0, 2, 2, 0, 1};
+        long[] predictions = new long[] {0, 0, 2, 2, 0, 2};
+        Operand<TInt64> result =
+            fops.math.confusionMatrix(tf.constant(labels), tf.constant(predictions));
+        long[][] expected =
+            new long[][] {
+              {2, 0, 0},
+              {0, 0, 1},
+              {1, 0, 2}
+            };
+        session.evaluate(tf.constant(expected), result);
+      }
+  }
+
+  @Test
+  public void testTensorDotValid() {
+    for (TestSession.Mode tfMode : tfModes)
+      try (TestSession session = TestSession.createTestSession(tfMode)) {
+        Ops tf = session.getTF();
+        FrameworkOps fops = FrameworkOps.create(tf);
+        int[] axes1 = new int[] {1, 2};
+        int[][] axes2 = new int[][] {{1}, {2}};
+        int[][] axes3 = new int[2][0];
+        int axes4 = 0;
+
+        Operand<TFloat32> a = tf.ones(tf.constant(Shape.of(3, 3)), TFloat32.class);
+        Operand<TFloat32> b = tf.constant(new float[][][] {{{2, 3, 1}}});
+
+        Operand<TFloat32> ans = fops.math.tensordot(a, b, axes1);
+        Operand<TFloat32> expected = tf.constant(new float[][][] {{{6}}, {{6}}, {{6}}});
+        session.evaluate(expected, ans);
+
+        ans = fops.math.tensordot(a, b, axes2);
+        expected = tf.constant(new float[][][] {{{6}}, {{6}}, {{6}}});
+        session.evaluate(expected, ans);
+
+        float[][][][][] expectedArray =
+            new float[][][][][] {
+              {{{{2, 3, 1}}}, {{{2, 3, 1}}}, {{{2, 3, 1}}}},
+              {{{{2, 3, 1}}}, {{{2, 3, 1}}}, {{{2, 3, 1}}}},
+              {{{{2, 3, 1}}}, {{{2, 3, 1}}}, {{{2, 3, 1}}}}
+            };
+        ans = fops.math.tensordot(a, b, axes3);
+        expected = tf.constant(expectedArray);
+        session.evaluate(expected, ans);
+
+        ans = fops.math.tensordot(a, b, axes4);
+        expected = tf.constant(expectedArray);
+        session.evaluate(expected, ans);
+      }
+  }
+
+  @Test
+  public void testTensorDotInValidAxis() {
+    for (TestSession.Mode tfMode : tfModes)
+      try (TestSession session = TestSession.createTestSession(tfMode)) {
+        Ops tf = session.getTF();
+        FrameworkOps fops = FrameworkOps.create(tf);
+        Operand<TFloat32> a = tf.constant(new float[][] {{1, 2}, {3, 4}});
+        Operand<TFloat32> b = tf.constant(new float[][] {{1, 2}, {3, 4}});
+        assertThrows(IllegalArgumentException.class, () -> fops.math.tensordot(a, b, -1));
+        assertThrows(IllegalArgumentException.class, () -> fops.math.tensordot(a, b, 3));
+        assertThrows(
+            IllegalArgumentException.class, () -> fops.math.tensordot(a, b, new int[] {1}));
+        assertThrows(
+            IllegalArgumentException.class, () -> fops.math.tensordot(a, b, new int[][] {{1}}));
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> fops.math.tensordot(a, b, new int[][] {{1}, {0, 1}}));
+
+        assertThrows(
+            ArrayIndexOutOfBoundsException.class,
+            () -> fops.math.tensordot(a, b, new int[][] {{0}, {7}}));
+      }
+  }
+
+  @Test
+  public void testReduceLogSumExp() {
+    for (TestSession.Mode tfMode : tfModes)
+      try (TestSession session = TestSession.createTestSession(tfMode)) {
+        Ops tf = session.getTF();
+        FrameworkOps fops = FrameworkOps.create(tf);
+        Operand<TFloat32> x =
+            tf.constant(
+                new float[][] {
+                  {0.43346116f, 0.8569728f, 0.57155997f, 0.0743812f, 0.63846475f},
+                  {0.8165283f, 0.26554802f, 0.37025765f, 0.8255019f, 0.45682374f},
+                  {0.93511814f, 0.52291054f, 0.80983895f, 0.11580781f, 0.8111686f},
+                  {0.49967498f, 0.27537802f, 0.48554695f, 0.28238368f, 0.7989301f},
+                  {0.8958915f, 0.84870094f, 0.56874424f, 0.08818512f, 0.13915819f}
+                });
+
+        Operand<TFloat32> result = fops.math.reduceLogSumExp(x, new int[] {0, 1}, false);
+        session.evaluate(3.7911222f, result);
+      }
+  }
+}
diff --git a/tensorflow-framework/src/test/java/org/tensorflow/framework/op/NnOpsTest.java b/tensorflow-framework/src/test/java/org/tensorflow/framework/op/NnOpsTest.java
new file mode 100644
index 00000000000..0436fdd57cf
--- /dev/null
+++ b/tensorflow-framework/src/test/java/org/tensorflow/framework/op/NnOpsTest.java
@@ -0,0 +1,68 @@
+package org.tensorflow.framework.op;
+
+import org.junit.jupiter.api.Test;
+import org.tensorflow.Operand;
+import org.tensorflow.framework.utils.TestSession;
+import org.tensorflow.op.Ops;
+import org.tensorflow.types.TFloat32;
+import org.tensorflow.types.TInt32;
+
+class NnOpsTest {
+  private final TestSession.Mode[] tfModes = {TestSession.Mode.EAGER, TestSession.Mode.GRAPH};
+
+  @Test
+  public void testSigmoidCrossEntropyWithLogits() {
+    for (TestSession.Mode tfMode : tfModes)
+      try (TestSession session = TestSession.createTestSession(tfMode)) {
+        Ops tf = session.getTF();
+        FrameworkOps fops = FrameworkOps.create(tf);
+        float[] x = new float[] {-100, -2, -2, 0, 2, 2, 2, 100};
+        float[] y = new float[] {0, 0, 1, 0, 0, 1, 0.5f, 1};
+
+        Operand<TFloat32> logits = tf.constant(x);
+        Operand<TFloat32> targets = tf.constant(y);
+        Operand<TFloat32> loss = fops.nn.sigmoidCrossEntropyWithLogits(targets, logits);
+        Operand<TFloat32> expected =
+            tf.constant(
+                new float[] {
+                  0.f, 0.126928f, 2.126928f, 0.6931472f,
+                  2.126928f, 0.126928f, 1.126928f, 0.f
+                });
+        session.evaluate(expected, loss);
+      }
+  }
+
+  @Test
+  public void testSoftmaxCrossEntropyWithLogits() {
+    for (TestSession.Mode tfMode : tfModes)
+      try (TestSession session = TestSession.createTestSession(tfMode)) {
+        Ops tf = session.getTF();
+        FrameworkOps fops = FrameworkOps.create(tf);
+        float[] x = new float[] {-100, -2, -2, 0, 2, 2, 2, 100};
+        float[] y = new float[] {0, 0, 1, 0, 0, 1, 0.5f, 1};
+
+        Operand<TFloat32> logits = tf.constant(x);
+        Operand<TFloat32> targets = tf.constant(y);
+        Operand<TFloat32> loss = fops.nn.softmaxCrossEntropyWithLogits(targets, logits, 0);
+
+        session.evaluate(249.0f, loss);
+      }
+  }
+
+  @Test
+  public void testSparseSoftmaxCrossEntropyWithLogits() {
+    for (TestSession.Mode tfMode : tfModes)
+      try (TestSession session = TestSession.createTestSession(tfMode)) {
+        Ops tf = session.getTF();
+        FrameworkOps fops = FrameworkOps.create(tf);
+        float[][] x = new float[][] {{0, 0}};
+        int[] y = new int[] {0};
+
+        Operand<TFloat32> logits = tf.constant(x);
+        Operand<TInt32> labels = tf.constant(y);
+        Operand<TFloat32> loss = fops.nn.sparseSoftmaxCrossEntropyWithLogits(labels, logits);
+
+        session.evaluate(0.69314718f, loss);
+      }
+  }
+}
diff --git a/tensorflow-framework/src/test/java/org/tensorflow/framework/op/SetOpsTest.java b/tensorflow-framework/src/test/java/org/tensorflow/framework/op/SetOpsTest.java
new file mode 100644
index 00000000000..41b2dfe4e4f
--- /dev/null
+++ b/tensorflow-framework/src/test/java/org/tensorflow/framework/op/SetOpsTest.java
@@ -0,0 +1,123 @@
+package org.tensorflow.framework.op;
+
+import static org.tensorflow.framework.utils.CastHelper.cast;
+
+import java.util.Arrays;
+import java.util.List;
+import org.junit.jupiter.api.Test;
+import org.tensorflow.Operand;
+import org.tensorflow.framework.utils.TestSession;
+import org.tensorflow.ndarray.Shape;
+import org.tensorflow.op.Ops;
+import org.tensorflow.types.TInt32;
+import org.tensorflow.types.TInt64;
+import org.tensorflow.types.TUint8;
+import org.tensorflow.types.family.TType;
+
+class SetOpsTest {
+
+  private final TestSession.Mode[] tfModes = {TestSession.Mode.EAGER, TestSession.Mode.GRAPH};
+
+  List<Class<? extends TType>> types = Arrays.asList(TInt32.class, TInt64.class, TUint8.class);
+
+  @Test
+  @SuppressWarnings({"unchecked", "rawtypes"})
+  public void testSetIntersectionMultirow2() {
+
+    for (TestSession.Mode tfMode : tfModes)
+      try (TestSession session = TestSession.createTestSession(tfMode)) {
+        Ops tf = session.getTF();
+        FrameworkOps fops = FrameworkOps.create(tf);
+        Operand<TInt32> a = tf.constant(new int[][] {{9, 1, 5}, {2, 4, 3}});
+        Operand<TInt32> b = tf.constant(new int[][] {{1, 9}, {1, 5}});
+        int[][] expected = new int[][] {{1, 9}, {0, 0}};
+        Shape expectedShape = Shape.of(2, 2);
+        for (Class<? extends TType> type : types) {
+          Operand aa = cast(tf, a, type);
+          Operand bb = cast(tf, b, type);
+          Operand<? extends TType> intersection = fops.sets.intersection(aa, bb);
+          session.evaluate(cast(tf, tf.constant(expected), type), intersection);
+          session.evaluate(tf.constant(expectedShape), tf.shape(intersection, TInt64.class));
+        }
+      }
+  }
+
+  @Test
+  @SuppressWarnings({"unchecked", "rawtypes"})
+  public void testSetIntersectionDuplicates2d() {
+
+    for (TestSession.Mode tfMode : tfModes)
+      try (TestSession session = TestSession.createTestSession(tfMode)) {
+        Ops tf = session.getTF();
+        FrameworkOps fops = FrameworkOps.create(tf);
+        Operand<TInt32> a = tf.constant(new int[][] {{1, 1, 3}});
+        Operand<TInt32> b = tf.constant(new int[][] {{1, 1}});
+        int[][] expected = {{1}};
+        Shape expectedShape = Shape.of(1, 1);
+        for (Class<? extends TType> type : types) {
+          Operand aa = cast(tf, a, type);
+          Operand bb = cast(tf, b, type);
+          Operand intersection = fops.sets.intersection(aa, bb);
+
+          session.evaluate(cast(tf, tf.constant(expected), type), intersection);
+
+          session.evaluate(tf.constant(expectedShape), tf.shape(intersection, TInt64.class));
+        }
+      }
+  }
+
+  @Test
+  @SuppressWarnings({"unchecked", "rawtypes"})
+  public void testDenseSetDifferenceMultirow2d() {
+
+    for (TestSession.Mode tfMode : tfModes)
+      try (TestSession session = TestSession.createTestSession(tfMode)) {
+        Ops tf = session.getTF();
+        FrameworkOps fops = FrameworkOps.create(tf);
+        Operand<TInt32> a = tf.constant(new int[][] {{1, 5, 9}, {4, 5, 3}});
+        Operand<TInt32> b = tf.constant(new int[][] {{1, 2, 6}, {1, 2, 2}});
+
+        for (Class<? extends TType> type : types) {
+          Operand aa = cast(tf, a, type);
+          Operand bb = cast(tf, b, type);
+          int[][] expected = {{5, 9, 0}, {3, 4, 5}};
+          // a- b
+          Shape expectedShape = Shape.of(2, 3);
+          Operand intersection = fops.sets.difference(aa, bb);
+          session.evaluate(cast(tf, tf.constant(expected), type), intersection);
+          session.evaluate(tf.constant(expectedShape), tf.shape(intersection, TInt64.class));
+
+          // b - a
+          expected = new int[][] {{2, 6}, {1, 2}};
+          expectedShape = Shape.of(2, 2);
+          intersection = fops.sets.difference(aa, bb, false);
+
+          session.evaluate(cast(tf, tf.constant(expected), type), intersection);
+          session.evaluate(tf.constant(expectedShape), tf.shape(intersection, TInt64.class));
+        }
+      }
+  }
+
+  @Test
+  @SuppressWarnings({"unchecked", "rawtypes"})
+  public void testDenseUnionMultirow2d() {
+
+    for (TestSession.Mode tfMode : tfModes)
+      try (TestSession session = TestSession.createTestSession(tfMode)) {
+        Ops tf = session.getTF();
+        FrameworkOps fops = FrameworkOps.create(tf);
+        Operand<TInt32> a = tf.constant(new int[][] {{9, 1, 5}, {2, 4, 3}});
+        Operand<TInt32> b = tf.constant(new int[][] {{1, 9}, {1, 2}});
+        int[][] expected = new int[][] {{5, 0}, {3, 4}};
+        for (Class<? extends TType> type : types) {
+          Operand aa = cast(tf, a, type);
+          Operand bb = cast(tf, b, type);
+          Shape expectedShape = Shape.of(2, 2);
+          // a- b
+          Operand intersection = fops.sets.difference(aa, bb);
+          session.evaluate(cast(tf, tf.constant(expected), type), intersection);
+          session.evaluate(tf.constant(expectedShape), tf.shape(intersection, TInt64.class));
+        }
+      }
+  }
+}
diff --git a/tensorflow-framework/src/test/java/org/tensorflow/framework/optimizers/GradientDescentTest.java b/tensorflow-framework/src/test/java/org/tensorflow/framework/optimizers/GradientDescentTest.java
index 17188499ee7..835ed8fdcaa 100644
--- a/tensorflow-framework/src/test/java/org/tensorflow/framework/optimizers/GradientDescentTest.java
+++ b/tensorflow-framework/src/test/java/org/tensorflow/framework/optimizers/GradientDescentTest.java
@@ -163,7 +163,7 @@ public void testDeterminism() {
           tf.withName("output").placeholder(TFloat32.class, Placeholder.shape(Shape.of(-1, 2)));
       Mean<TFloat32> loss =
           tf.math.mean(
-              tf.nn.raw.softmaxCrossEntropyWithLogits(output, placeholder).loss(), tf.constant(0));
+              tf.nn.softmaxCrossEntropyWithLogits(output, placeholder).loss(), tf.constant(0));
       lossName = loss.op().name();
 
       GradientDescent gd = new GradientDescent(g, 10.0f);