std::vector<int64_t> in_shape{input_type.getShape().vec()};
  std::vector<int64_t> out_shape{output_type.getShape().vec()};

  // If the reshape changes the number of dimensions so it cannot be interpreted
  // as a transpose.
  if (in_shape.size() != out_shape.size()) {
    return false;
  }

  in_shape.erase(std::remove(in_shape.begin(), in_shape.end(), 1),
                 in_shape.end());

  ArrayRef<int64_t> in_shape = ranked_type.getShape();
  if (in_shape.empty() || in_shape[0] < 0) {
    return context_op->emitOpError()
           << "A map_outside_compilation op's input and output shapes must "
              "have rank at least one and the first dimension must be known.";
  }
  int64_t split_size = in_shape[0] / num_cores_per_replica;
  if (in_shape[0] % num_cores_per_replica != 0) {

  auto y = Sub(scope_, x1, x2);
  RunTest({x1, x2}, {x1_shape, x2_shape}, {y}, {x1_shape});
}

TEST_F(NaryGradTest, Mul) {
  TensorShape x1_shape({3, 2, 5});
  TensorShape x2_shape({2, 5});
  auto x1 = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x1_shape));
  auto x2 = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x2_shape));
  auto y = Mul(scope_, x1, x2);

  // https://www.tensorflow.org/api_docs/cc/class/tensorflow/ops/batch-mat-mul
  int64_t x_row_dim = x_shape[x_shape.size() - 2];
  int64_t x_col_dim = x_shape[x_shape.size() - 1];
  int64_t y_row_dim = y_shape[y_shape.size() - 2];
  int64_t y_col_dim = y_shape[y_shape.size() - 1];
  int64_t out_row_dim = output_shape[output_shape.size() - 2];
  int64_t out_col_dim = output_shape[output_shape.size() - 1];

          contracting_dims.add(c)
      x_signature = [
          None if c not in contracting_dims else x_shape[cidx]
          for cidx, c in enumerate(x_labels)
      ]
      y_signature = [
          None if c not in contracting_dims else y_shape[cidx]
          for cidx, c in enumerate(y_labels)
      ]
    return x_shape, y_shape, bias_shape, x_signature, y_signature

  def _create_einsum_model(
      self,

      x_signature = [
          None if c not in contracting_dims else x_shape[cidx]
          for cidx, c in enumerate(x_labels)
      ]
      y_signature = [
          None if c not in contracting_dims else y_shape[cidx]
          for cidx, c in enumerate(y_labels)
      ]
    return x_shape, y_shape, bias_shape, x_signature, y_signature

  def _create_einsum_model(
      self,

    Value result = reduction.getResult(0);

    // The mean op needs to divide by the product of the reduced dimensions.
    if (std::is_same<OpTy, TF::MeanOp>::value) {
      Value in_shape = rewriter.create<shape::ShapeOfOp>(loc, op.getInput());
      Value divisor_count = rewriter.create<arith::ConstantIndexOp>(loc, 1);
      for (size_t i = 0; i < input_shape.size(); ++i) {

  // CHECK-DAG: %[[RHS:.+]] = mhlo.subtract %[[OFFSET]], %[[MUL_MEAN]] : tensor<?xf32>
  // CHECK-DAG: %[[X_SHAPE:.+]] = shape.shape_of %[[X]] : tensor<?x?x?x?xf32> -> tensor<4xindex>
  // CHECK-DAG: %[[MULTIPLIER_BCAST:.+]] = "mhlo.dynamic_broadcast_in_dim"(%[[MULTIPLIER]], %[[X_SHAPE]]) <{broadcast_dimensions = dense<1> : tensor<1xi64>}> : (tensor<?xf32>, tensor<4xindex>) -> tensor<?x?x?x?xf32>

#include "tensorflow/c/experimental/saved_model/public/signature_def_param_list.h"
#include "tensorflow/c/experimental/saved_model/public/tensor_spec.h"
#include "tensorflow/c/tf_datatype.h"
#include "tensorflow/c/tf_shape.h"
#include "tensorflow/c/tf_status.h"
#include "tensorflow/c/tf_tensor.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/stringpiece.h"

  // TFR_ConstantTensorOp (
  //   ConstantOp (ConstAttr<F32Attr (in_scale[0] * in_scale[1] /
  //   out_scale))
  // )
  // Currently, all decompositions using this pattern (Conv2D, FC) have the
  // following preconditions:
  // * out_scale: float scalar attribute
  // * in_scale[0] (input scale): float scalar, given by tf.Const -> tfr.cast
  // * in_scale[1] (filter scale): float scalar/vector