if (BroadcastVector<double>(shaped.getDimSize(quant_dim), scales) ||
        BroadcastVector<int64_t>(shaped.getDimSize(quant_dim), zero_points)) {
      return {};
    }
  } else if ((new_shape.size() == shape.size() + 1) && new_shape.front() == 1) {
    // Handle the [A, B, C] -> [1, A, B, C] reshape case.
    if (!(std::equal(shape.begin(), shape.end(), new_shape.begin() + 1) &&

namespace tensorflow {

// Attribute marking output tensor shapes inferred by XLA. Attribute value is
// a list of PartialTensorShape objects.
extern const char kXlaInferredShapesAttrName[];

// Infers output shapes for all nodes in graph `g`. The output shapes will be
// stored in node attribute `kXlaInferredShapesAttrName`.
//
// We have to perform shape inference before encapsulation because after

  xla_tensor->WaitForDefinitionEventOnStream(device_to_host_stream.get());

  // Transfer manager requires the shape of the shaped buffer to be the same as
  // literal shape except for the layout.  Set the literal to use xla_tensor's
  // shape as it is derived from the cpu_tensor's shape using
  // shape_representation_fn_.
  xla::MutableBorrowingLiteral literal;
  TF_CHECK_OK(HostTensorToMutableBorrowingLiteral(

  pm.addNestedPass<FuncOp>(
      mlir::TF::CreateDropWhileShapeInvariantInDeviceClusterPass());
  // Run another shape inference pass because resource decomposition might have
  // created new partial types. Also, after dropping `shape_invariant` attribute
  // from While/WhileRegion ops within cluster would lead to more precise
  // shapes.
  pm.addPass(mlir::TF::CreateTFShapeInferencePass());
  pm.addNestedPass<FuncOp>(mlir::createCanonicalizerPass());

              fused_batch_norm_op, [&](::mlir::Diagnostic &diag) {
                diag << "Shapes of mean and variance should be 1D and "
                        "compatible with x";
              });
        }
      }

      // Check if output shape and input shape are compatible.
      auto x_type = (*x.begin()).getType();
      auto y_type = (*root.getODSResults(0).begin()).getType();

// true when the args are not equal.
struct SignatureNotEqual {
  bool operator()(const Tensor& arg, const Tensor& other) {
    return arg.dtype() != other.dtype() || arg.shape() != other.shape() ||
           arg.tensor_data() != other.tensor_data();
  }
  bool operator()(const TensorTypeAndShape& arg,
                  const TensorTypeAndShape& other) {

class IdentityNOp;

// Returns the RankedTensorType for the given operand. TensorFlow constant ops
// may have non-static shape because the shape is not propagated during constant
// folding. If the defining op for the given operand is a constant op, this
// routine uses the constant op's attribute to get the actual shape.
RankedTensorType GetRankedTensorTypeForOperand(Value operand) {
  DenseElementsAttr attr;

    }

    std::vector<PartialTensorShape> shapes;
    if (!GetNodeAttr(e->src()->attrs(), kXlaInferredShapesAttrName, &shapes)
             .ok()) {
      return std::nullopt;
    }

    const PartialTensorShape shape = shapes[e->src_output()];
    if (!shape.IsFullyDefined()) {
      return std::nullopt;
    }

    results[e->dst_input()] = shape;
  }
  return results;
}

// Given a tensor of `shape` and `type`, as what shape should it be stored on
// the TPU device? This function tranposes or flattens the excessively-padded
// tensors to rank 1, but leaves other tensor shapes alone.
absl::StatusOr<xla::Shape> TpuShapeRepresentation(
    const TensorShape& shape, DataType type, bool use_fast_memory,
    XlaLayoutPreference layout_preference) {
  xla::Shape xla_shape;
  TF_RETURN_IF_ERROR(

    if (ShapedType::isDynamic(num_index_dims)) return failure();

    auto updates = op.getUpdates();

    // Broadcast scalar `updates` in into expected shape as following shape:
    // updates.shape == indices.shape[:-1] + tensor.shape[indices.shape[-1]:]
    if (updates_ty.getRank() == 0 &&
        (std::is_same<OpTy, TF::TensorScatterUpdateOp>::value ||
         std::is_same<OpTy, TF::TensorScatterAddOp>::value)) {