llvm::dyn_cast<mlir::TF::TPUPartitionedOutputV2Op>(user))
    return NullUnlessSharded(partitioned_output);

  if (auto assign_var = llvm::dyn_cast<mlir::TF::AssignVariableOp>(user))
    if (auto partitioned_input =
            assign_var.getResource()
                .getDefiningOp<mlir::TF::TPUPartitionedInputV2Op>())
      return NullUnlessSharded(partitioned_input);

  return nullptr;
}