continue;

      // Check that there is only a single global tensor associated with arg.
      const TF::ResourceDataflowState *latticeElement =
          solver.lookupState<TF::ResourceDataflowState>(val);
      if (!latticeElement || latticeElement->getValue().ops.size() != 1)
        continue;

      // Don't freeze mutable tensors.

  module->walk([&](TF::TPUCopyWithDynamicShapeOp op) {
    const Device *state;
    for (auto result : op->getResults()) {
      state = solver.lookupState<Device>(result);
      if (state) break;
    }
    if (!state || !state->hasDevice()) {
      return WalkResult::advance();
    }
    op->setAttr(kDevice, state->getDevice());

  }

  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) {
    // Can't actually happen. Even for gaps in the dataflow, we'll receive
    // the initial state.
    return;
  }
  auto ops = state->getValue().ops;

    resource = read.getResource();
  } else if (auto write = llvm::dyn_cast<TF::AssignVariableOp>(op)) {
    resource = write.getResource();
  }
  const TF::ResourceDataflowState* state =
      solver.lookupState<TF::ResourceDataflowState>(resource);
  if (!state) {
    return nullptr;
  }
  auto ops = state->getValue().ops;
  if (ops.size() != 1) {
    return nullptr;
  }
  Operation* source = *ops.begin();

      (void)id;
    }
  }

  for (auto& arg : op->getOpOperands()) {
    if (!IsResourceType(arg.get().getType())) {
      continue;
    }
    auto lattice =
        solver.lookupState<TF::IsCompositeDataflowState>(arg.get())->getValue();
    bool is_read = read_array.contains(arg.getOperandNumber());
    bool is_update = update_array.contains(arg.getOperandNumber());