nodes_[inst] = node;
  return absl::OkStatus();
}

bool IsEntryFunctionArg(BlockArgument arg) {
  return arg.getParentRegion()->getParentOfType<FuncOp>().getName() == "main";
}

// Creates argument nodes from Block argument. If a name is supplied, that
// name will be used instead of generating a unique name.
Status Exporter::AddArgumentNode(BlockArgument arg, unsigned index,

            /*quant_dim=*/-1, num_bits, narrow_range, is_signed);
        builder.setInsertionPoint(block, insertion_point);
        auto q_op = builder.create<quantfork::QuantizeCastOp>(
            loc, params.getValue(), arg);
        auto dq_op = builder.create<quantfork::DequantizeCastOp>(
            loc, input_type, q_op.getResult());
        arg.replaceAllUsesWith(dq_op.getResult());
        q_op.setOperand(arg);
      }

        // Keep track of IDs of resources that are allocated by ops with
        // `UniqueResourceAllocation` trait, this can be utilized for while-loop
        // parallelization (every iteration creates a new unique resource).
        if (op->hasTrait<OpTrait::TF::UniqueResourceAllocation>()) {
          unique_resource_allocation_ids_.insert(resource_handle.id);
        }
      }

    const llvm::StringMap<Attribute>& derived_attrs) const {
  // Create the new op
  Location loc = call_op.getLoc();
  rewriter.setInsertionPointAfter(call_op);
  std::string tf_op_name = GetTFOpName(call_op.getCallee());
  OperationState new_state(loc, tf_op_name, inputs, output_types, attr_list);
  Operation* new_op = rewriter.create(new_state);
  if (materialize_derived_attrs_) {

  auto launch = builder->create<tf_device::LaunchOp>(
      before_op->getLoc(), builder->getStringAttr(host_device),
      launch_result_types);
  launch.getBody().push_back(launch_block);

  builder->setInsertionPointToEnd(&launch.GetBody());
  builder->create<tf_device::ReturnOp>(before_op->getLoc(), launch_results);

  return launch;
}

// Creates a new_launch after `before_op` with `ops_to_move` from

    if (failed(applyPatternsAndFoldGreedily(func, frozen_patterns))) {
      func.emitError() << "quant-propagate-quantize-type failed.";
      signalPassFailure();
    }
  }
}

}  // namespace

// Creates an instance of the TensorFlow dialect PropagateQuantizeType pass.
std::unique_ptr<OperationPass<ModuleOp>> CreatePropagateQuantizeTypePass() {
  return std::make_unique<PropagateQuantizeType>();
}

  // supported.
  const auto& device_list =
      devices.find(tensorflow::GetDeviceAliasForLogicalCore(0))->getSecond();

  // Create the state variable for each device.
  for (llvm::StringRef device : device_list) {
    state_vars.push_back(builder->create<TF::VarHandleOp>(
        loc,
        llvm::ArrayRef<Type>{RankedTensorType::get(
            {}, TF::ResourceType::get(llvm::ArrayRef<TensorType>{key_type},

    auto static_shape_attr =
        mlir::DenseIntElementsAttr::get(static_shape_type, static_shape);
    return rewriter.create<TF::ConstOp>(loc, static_shape_attr).getOutput();
  }

  // If the shape is not static, create a new ShapeOp.
  BoolAttr false_attr = rewriter.getBoolAttr(false);
  return rewriter
      .create<TF::ShapeOp>(loc, input,
                           /*use_32bit=*/false_attr)
      .getOutput();
}

  auto out_segids_cst = builder.create<TF::ConstOp>(
      builder.getI32TensorAttr(flattened_out_segids));
  auto contracting_segids_cst = builder.create<TF::ConstOp>(
      builder.getI32TensorAttr(flattened_contracting_segids));
  auto num_segids_tensor =
      builder.create<TF::ConstOp>(builder.getI32IntegerAttr(1));
  auto flattened_out_dims = builder.create<TF::UnsortedSegmentProdOp>(

        return false;
      }
    }
    rewriter.setInsertionPointAfter(op);
    auto q = rewriter.create<quantfork::QuantizeCastOp>(op->getLoc(), cast_type,
                                                        op.getResult());
    auto dq = rewriter.create<quantfork::DequantizeCastOp>(op->getLoc(),
                                                           expressed_type, q);