device_ordinal);
  for (auto& index_to_buffer : shaped_buffer.buffers()) {
    xla::Shape subshape =
        xla::ShapeUtil::GetSubshape(on_device_shape, index_to_buffer.first);
    uint64 size =
        client->backend().transfer_manager()->GetByteSizeRequirement(subshape);
    TF_ASSIGN_OR_RETURN(se::OwningDeviceMemory buffer,
                        client->backend().memory_allocator()->Allocate(

    sq2 = gen_composite_ops.my_add_n([t1, t2])
    sq3 = gen_composite_ops.my_add_n([t1, t2, t3])
    self.assertAllEqual(sq1.numpy().reshape(-1), [1, 2, 3, 4])
    self.assertAllEqual(sq2.numpy().reshape(-1), [2, 4, 6, 8])
    self.assertAllEqual(sq3.numpy().reshape(-1), [3, 6, 9, 12])

  def testBiasedDense(self):
    t1 = constant_op.constant([[1.0, 2.0], [3.0, 4.0]])
    t2 = constant_op.constant([[1.0, 2.0], [3.0, 4.0]])

// CHECK-NEXT: %[[reshape:.*]] = "tfl.reshape"(%[[conv]], %[[shape]]) : (tensor<1x112x112x32x!quant.uniform<i8:f32, 0.023528476789885875>>, tensor<2xi32>) -> tensor<1x401408x!quant.uniform<i8:f32, 0.023528476789885875>>

  // CHECK: %[[v1:.*]] = "tf.Reshape"(%arg1, %[[cst_1]]) : (tensor<5x2xf32>, tensor<3xi64>) -> tensor<5x2x1xf32>
  // CHECK: %[[v2:.*]] = "tf.BatchMatMulV2"(%[[v0]], %[[v1]]) <{adj_x = false, adj_y = false}> : (tensor<5x7x2xf32>, tensor<5x2x1xf32>) -> tensor<5x7x1xf32>
  // CHECK: %[[v3:.*]] = "tf.Reshape"(%[[v2]], %[[cst_2]]) : (tensor<5x7x1xf32>, tensor<2xi64>) -> tensor<5x7xf32>

// CHECK: %[[slice:.*]] = "tf.Slice"(%arg0, %[[tensor_scatter_update]], %[[arg2_i64]]) : (tensor<?x2xf32>, tensor<2xi64>, tensor<2xi64>) -> tensor<*xf32>
// CHECK: %[[reshape:.*]] = "tf.Reshape"(%[[slice]], %[[cst_1]]) : (tensor<*xf32>, tensor<1xi64>) -> tensor<*xf32>
// CHECK: return %[[reshape]] : tensor<*xf32>

// -----

// Tests that the converted `tf.Slice` has the correct number of dimensions

        DenseElementsAttr::get(shape_spec_type, new_shape.getShape());
    rewriter.setInsertionPointAfter(weight_op);
    auto new_shape_const = rewriter.create<TF::ConstOp>(
        weight_op->getLoc(), shape_spec_type, new_shape_const_attr);
    auto reshape_op = rewriter.create<TF::ReshapeOp>(
        weight_op->getLoc(), new_shape, weight_op->getResult(0),
        new_shape_const);

    // Get the unbroadcasted shapes in the operand order.
    std::array<llvm::ArrayRef<int64_t>, 2> operand_shapes;
    operand_shapes[i] = broadcast_arg_type.getShape();
    operand_shapes[1 - i] = argument_type.getShape();

    // Check that the input of the broadcast and the other operand is broadcast
    // compatible.
    llvm::SmallVector<int64_t, 4> broadcasted_shape;

// identity dimensions permuted.
def IsTransposeTrivial : Constraint<CPred<
  "TFL::IsTransposeTrivial($0.getType().cast<ShapedType>().getShape(), $1)">>;

// Constraint that checks if the reshape op is equivalent to a transpose op.
// This is true if the reshape op is a trivial reshape op, meaning no change in
// the order of non-identity dimensions.
def IsReshapeEquivalentToTranspose : Constraint<CPred<

  if (!type.hasRank()) return type;

  auto shape = type.getShape();
  SmallVector<int64_t, 4> new_shape(shape.size());

  for (int i = 0; i < permutation.size(); ++i) {
    int64_t index = permutation[i];
    assert(index < shape.size());
    new_shape[index] = shape[i];
  }

  return type.clone(new_shape);
}

// Move Transpose operations that permute `op` operands after the `op`.

  auto shape = mlir::cast<ShapedType>(value.getType());
  if (shape.getRank() != 1) {
    SmallVector<int64_t> new_shape;
    new_shape.push_back(shape.getNumElements());
    value = builder.create<TF::ReshapeOp>(
        loc, value, Create1DConstValue(builder, loc, new_shape));
  }
  return ConstantFoldOpIfPossible(value.getDefiningOp()).front();
}