return failure();
    }

    // Compute multiplier = scale / sqrt(variance + epsilon)
    Value multiplier = rewriter.create<mhlo::AddOp>(
        bn_op.getLoc(), bn_op.getVariance(), epsilon);
    multiplier = rewriter.create<mhlo::RsqrtOp>(bn_op.getLoc(), multiplier);
    multiplier = rewriter.create<mhlo::MulOp>(bn_op.getLoc(), multiplier,
                                              bn_op.getScale());

      absl::StrCat("device: ", DEVICE_TPU_SYSTEM, " device"));
  devices->push_back(std::make_unique<VirtualDevice>(options.env, attrs));
  VLOG(1) << "Created TPU_SYSTEM device. This host has " << device_count
          << " TPUs";

  return absl::OkStatus();
}

}  // namespace

void RegisterTpuDeviceToDeviceCopy() {
  static auto* const register_tpu_tpu_copy = new CopyTensor::Registration(

    if (tpu_replicated_input.getIsMirroredVariable()) {
      mirrored_variable_indices.push_back(pos_and_input.index());
    }
  }

  // Create replicate op.
  auto result_types = GetClusterResultTypes(cluster, partitioned_outputs);
  auto replicate_op = builder.create<mlir::tf_device::ReplicateOp>(
      cluster.getLoc(), num_replicas,
      llvm::SmallDenseMap<llvm::StringRef,

// function, which might not exist in case of multi-signature graphs. In that
// case, this pass will create a new main function, which calls signature
// functions.
//
// An already existing @main function will be renamed by attaching a numeric
// suffix like `@main_0` to avoid conflict with the newly created main function.
class InsertMainFunctionPass
    : public PassWrapper<InsertMainFunctionPass, OperationPass<ModuleOp>> {

    return builder_.create<mlir::tf_executor::SwitchOp>(loc, types, operands,
                                                        result.attributes);
  }
  if (node.IsMerge()) {
    return builder_.create<mlir::tf_executor::MergeOp>(loc, types, operands,
                                                       result.attributes);
  }
  if (node.IsNextIteration()) {

    // Finally, use the quantization parameter to create the quantize and
    // dequantize ops, and insert them between the tf.FakeQuantWithMinMaxVarsOp
    // and its users.
    Value value = tf_op.getOutputs();
    auto quantize = rewriter.create<quantfork::QuantizeCastOp>(
        tf_op.getLoc(), qtype.getValue(), value);
    auto dequantize = rewriter.create<quantfork::DequantizeCastOp>(

    auto new_func_type = builder.getFunctionType(block.getArgumentTypes(),
                                                 terminator->getOperandTypes());
    func.setType(new_func_type);
  });
}

// Creates an instance of the pass to decompose the TF ops.
std::unique_ptr<OperationPass<ModuleOp>> CreateRewriteQuantizedIOPass() {
  return std::make_unique<RewriteQuantizedIOPass>();
}

      ToNodeProperties(R"pb(
        name: 'XTimesY' op: 'XTimesY' input: 'a' input: 'b'
      )pb");
  (*(callsite->node_def.mutable_attr()))["_XlaMustCompile"] = BoolAttr(true);

  // Note: need to set attribute on the created node.
  Status status = xla_kernel_creator.CreateKernel(flr_, callsite, &kernel_);
  ASSERT_TRUE(status.ok()) << status.ToString();

  EXPECT_EQ("XTimesY", kernel_->name());

// Test to ensure that a registered graph optimization pass is only executed
// once (i.e., on the main function side) in running distributed functions.
// This test creates a cluster with two workers, create a variable on the
// second worker, and run a distributed function (VariableAddFunction) whose ops
// span the local and remote workers. If the graph optimization pass is executed

  auto graph_op = builder.create<tf_executor::GraphOp>(
      loc, func.getFunctionType().getResults());
  graph_op.getBody().push_back(new Block);
  builder.setInsertionPointToEnd(&graph_op.GetBody());
  auto island = builder.create<tf_executor::IslandOp>(
      loc, func.getFunctionType().getResults(),
      tf_executor::ControlType::get(&getContext()), ArrayRef<Value>());
  // Create Fetch.