auto type = result.getType();
  result.setType(converter.convertType(result.getType()));
  auto new_value = converter.materializeArgumentConversion(
      rewriter, result.getLoc(), type, {result});
  rewriter.replaceAllUsesExcept(result, new_value, new_value.getDefiningOp());
}

// Wrap operands in an an unrealized cast to create a cast to buffer any type
// changes to the operand, and apply type converter to operands:

      Operation *aggregator_op = rewriter.create<TF::CustomAggregatorOp>(
          op->getLoc(), output_types, value, attributes);

      Value aggregator_op_result = aggregator_op->getOpResult(0);
      value.replaceAllUsesExcept(aggregator_op_result, aggregator_op);
    }

    return success();
  }

 private:
  CalibrationOptions calib_opts_;
};

void InsertCustomAggregationOpsPass::runOnOperation() {

        // 'sharding' attribute.
        auto updated_arg = builder.create<TF::XlaShardingOp>(
            func.getLoc(), arg.getType(), arg, sharding, sharding);
        func.getArgument(i).replaceAllUsesExcept(
            updated_arg, llvm::SmallPtrSet<Operation*, 1>({updated_arg}));
      }

      func.removeArgAttr(i, builder.getStringAttr(kShardingAttr));
    }
  }
}

  for (auto [original_output_value, xla_call_module_op_result_value] :
       llvm::zip_equal(outputs, xla_call_module_op->getResults())) {
    original_output_value.replaceAllUsesExcept(xla_call_module_op_result_value,
                                               /*exceptedUser=*/nullptr);
  }
}

// Replaces the StableHLO ops with a separate XlaCallModuleOp, then wires it

  Operation* terminator_op =
      old_op->getParentOfType<func::FuncOp>().front().getTerminator();
  llvm::SmallPtrSet<mlir::Operation*, 1> exceptions = {terminator_op};
  old_val.replaceAllUsesExcept(new_val, exceptions);
}

struct TensorArrayStats {
  // Whether a write op should accumulate with the old value. Set to true if
  // this is a gradient.
  bool accumulate_on_write;