Region &body = cluster.getBody();
      visitUsedValuesDefinedAbove(body, [&](OpOperand *use) {
        Value constant = use->get();
        auto const_op = dyn_cast_or_null<TF::ConstOp>(constant.getDefiningOp());
        if (!const_op) return;

        // Filter constants using user provided predicate function.
        if (filter && !filter(cluster, const_op.getValue())) return;

// returns the owner of the Block.
Operation* GetOpOfValue(Value value) {
  if (auto block_arg = mlir::dyn_cast<BlockArgument>(value))
    return block_arg.getOwner()->getParentOp();

  return value.getDefiningOp();
}

// TODO(b/158596585): Replace this with a cost model analysis.
bool IsTrivialUnaryOperation(Operation* op) {
  return llvm::isa<TF::CastOp, TF::IdentityOp>(op);
}

def F32ElementsAttr : ElementsAttrBase<
    CPred<"$_self.cast<ElementsAttr>().getShapedType().getElementType().isF32()">, "float constant tensor">;
def DefinedByConv2D : Constraint<CPred<"llvm::isa_and_nonnull<mlir::TF::Conv2DOp>($0.getDefiningOp())">>;
// Checks if the value has only one user.
def HasOneUse : Constraint<CPred<"$0.hasOneUse()">>;

// If we see a Conv2D op followed by Mul, then multiply the filter
// with the value in Mul.
def FuseMulAndConv2D :

  // Broadcast bias value if unmatched with output shape.
  auto bcast_op = TryCast<BroadcastInDimOp>(bias_op.getDefiningOp(),
                                            /*name=*/"broadcast_in_dim_op");

  if (failed(bcast_op)) {
    bcast_op = TryCast<DynamicBroadcastInDimOp>(
        bias_op.getDefiningOp(),
        /*name=*/"dynamic_broadcast_in_dim_op");
  }
  // Update the bias type for both static and dynamic broadcasts.

mlir::Operation* GetXlaShardingFromOperand(Value value) {
  Value value_to_visit = value;
  if (auto read_var = value_to_visit.getDefiningOp<mlir::TF::ReadVariableOp>())
    value_to_visit = read_var.getResource();

  if (auto partitioned_input =
          value_to_visit.getDefiningOp<mlir::TF::TPUPartitionedInputV2Op>()) {
    return NullUnlessSharded(partitioned_input);
  }

  return nullptr;
}

    // Only values defined by IslandOp were inserted in the worklist.
    IslandOp current_island = cast<IslandOp>(control_barrier.getDefiningOp());

    for (auto control_input : current_island.getControlInputs()) {
      ops_connected_to_fetch->insert(control_input.getDefiningOp());
      if (IsNoOpControlBarrier(control_input))
        control_barrier_worklist.push(control_input);
    }
    current_island.erase();

    resource = write.getResource();
  } else if (auto next = llvm::dyn_cast<TF::IteratorGetNextOp>(op)) {
    resource = next.getIterator();
  }

  if (llvm::dyn_cast_or_null<TF::VarHandleOp>(resource.getDefiningOp())) {
    return;  // We're already directly after a VarHandleOp.
  }

  const TF::ResourceDataflowState* state =
      solver.lookupState<TF::ResourceDataflowState>(resource);
  if (!state) {

        op_result.setType(ToBfloat16Type(original_type));
        const Value cast =
            builder.create<TF::CastOp>(op->getLoc(), original_type, op_result);
        op_result.replaceAllUsesExcept(cast, cast.getDefiningOp());
      }
    }
    return WalkResult::advance();
  });

  if (result.wasInterrupted()) return signalPassFailure();
}

}  // namespace

    if (handle_user == assign_var) continue;
    auto assign_var_user = dyn_cast<TF::AssignVariableOp>(handle_user);
    if (!assign_var_user) continue;
    if (assign_var_user.getValue().getDefiningOp() == cluster_func)
      return resource;
  }

  resource.resource = assign_var.getResource();
  resource.subtype = assign_var.getValue().getType();
  return resource;
}

  // hoisted.
  auto read_var_op = llvm::dyn_cast<ReadVariableOp>(op);
  if (!read_var_op) return false;
  auto var_handle_op = llvm::dyn_cast_or_null<VarHandleOp>(
      read_var_op.getResource().getDefiningOp());
  if (!var_handle_op) return false;
  return read_only_vars.contains(GetResourceHandle(var_handle_op));
}

bool ShouldMoveOutOfRegion(
    Operation *op, Region *region,