Home
last modified time | relevance | path

Searched defs:dtensor_recv (Results 1 – 3 of 3) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/expansions/
H A Ddtensor_op_spmd_expander.cc60 mlir::TF::DTensorRecv dtensor_recv) { in ValidateSendRecvLayoutConfiguration()
273 auto dtensor_recv = llvm::dyn_cast<mlir::TF::DTensorRecv>(recv_op); in ExpandOp() local
390 auto dtensor_recv = llvm::cast<mlir::TF::DTensorRecv>(op); in ExpandOp() local
531 mlir::TF::DTensorRecv dtensor_recv = in ComputeLayoutForward() local
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/
H A Ddtensor_send_recv.cc142 mlir::TF::DTensorRecv dtensor_recv) { in LowerDTensorRecvToXlaOp()
150 mlir::TF::DTensorRecv dtensor_recv, mlir::Type output_type) { in LowerDTensorRecvToXlaOp()
228 const Mesh& send_mesh, mlir::TF::DTensorRecv dtensor_recv) { in LowerDTensorRecvFromCPUToTFOp()
H A Dlower_send_recv.cc75 auto dtensor_recv = llvm::dyn_cast<mlir::TF::DTensorRecv>(*recv_op); in LowerDTensorSendRecvsOps() local