Home
last modified time | relevance | path

Searched refs:InferSPMDExpandedLocalShape (Results 1 – 25 of 34) sorted by relevance

12

/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/expansions/
H A Dconv_spmd_expander.cc152 return InferSPMDExpandedLocalShape(conv_op); in HandleConv()
285 return InferSPMDExpandedLocalShape(conv_op); in HandleConv()
305 return InferSPMDExpandedLocalShape(conv_op); in HandleConvBackpropInput()
338 return InferSPMDExpandedLocalShape(conv_op); in HandleConvBackpropFilter()
360 return InferSPMDExpandedLocalShape(max_pool_grad_op); in HandleMaxPoolGradOp()
413 return InferSPMDExpandedLocalShape(op); in ExpandOp()
H A Dmeta_spmd_expander.cc162 return InferSPMDExpandedLocalShape(op); in ExpandOp()
218 return InferSPMDExpandedLocalShape(op); in ExpandOp()
286 return InferSPMDExpandedLocalShape(op); in ExpandOp()
382 return InferSPMDExpandedLocalShape(op); in ExpandOp()
672 return InferSPMDExpandedLocalShape(op); in ExpandOp()
901 return InferSPMDExpandedLocalShape(op); in ExpandOp()
1080 return InferSPMDExpandedLocalShape(one_hot_op); in ExpandOp()
1156 InferSPMDExpandedLocalShape(op); in ExpandOp()
H A Dresource_spmd_expander.cc123 return InferSPMDExpandedLocalShape(op); in ExpandOp()
138 InferSPMDExpandedLocalShape(op); in ExpandOp()
215 return InferSPMDExpandedLocalShape(op); in ExpandOp()
H A Dtop_k_spmd_expander.cc67 new_op = InferSPMDExpandedLocalShape(new_op); in ExpandOp()
75 return InferSPMDExpandedLocalShape(op); in ExpandOp()
H A Dfill_spmd_expander.cc55 return InferSPMDExpandedLocalShape(op); in ExpandOp()
89 return InferSPMDExpandedLocalShape(new_fill); in ExpandOp()
H A Dsplit_spmd_expander.cc97 return InferSPMDExpandedLocalShape(op); in ExpandOp()
150 return InferSPMDExpandedLocalShape(op); in ExpandOp()
H A Dslice_spmd_expander.cc532 new_op = InferSPMDExpandedLocalShape(new_op); in ExpandOp()
612 op = InferSPMDExpandedLocalShape(op); in ExpandOp()
760 op = InferSPMDExpandedLocalShape(op); in ExpandOp()
910 op = InferSPMDExpandedLocalShape(op); in ExpandOp()
H A Ddataparallel_spmd_expander.cc231 op = InferSPMDExpandedLocalShape(op); in RelayoutOperandsAndOutputs()
301 return InferSPMDExpandedLocalShape(op); in ExpandOp()
H A Drange_spmd_expander.cc36 return InferSPMDExpandedLocalShape(op); in ExpandOp()
H A Dtensorlist_getitem_spmd_expander.cc41 return InferSPMDExpandedLocalShape(output.getDefiningOp()); in ExpandOp()
H A Didentity_n_spmd_expander.cc72 return InferSPMDExpandedLocalShape(identity_op.getOperation()); in ExpandOp()
H A Dtrivial_spmd_expander.h29 return InferSPMDExpandedLocalShape(op); in ExpandOp()
H A Dtrivial_spmd_expander.cc76 return InferSPMDExpandedLocalShape(op); in ExpandOp()
H A Dqr_spmd_expander.cc59 op = InferSPMDExpandedLocalShape(op); in ExpandOp()
H A Dgather_spmd_expander.cc164 return InferSPMDExpandedLocalShape(new_gather); in ExpandOp()
456 InferSPMDExpandedLocalShape(new_gather); in ExpandOp()
H A Dcumsum_spmd_expander.cc88 op = InferSPMDExpandedLocalShape(op); in ExpandOp()
H A Delementwise_spmd_expander.cc117 return InferSPMDExpandedLocalShape(op); in ExpandOp()
H A Dnullary_spmd_expander.cc79 return InferSPMDExpandedLocalShape(op); in ExpandOp()
H A Dsegmentation_spmd_expander.cc122 InferSPMDExpandedLocalShape(new_sum_op); in ExpandOp()
H A Dreplicated_spmd_expander.cc59 op = InferSPMDExpandedLocalShape(op); in ReplicatedRelayoutOperandsAndOutputs()
H A Dconcat_spmd_expander.cc94 return InferSPMDExpandedLocalShape(op); in ExpandOp()
H A Dexpanddims_spmd_expander.cc38 InferSPMDExpandedLocalShape(op); in ExpandOp()
H A Dsqueeze_spmd_expander.cc147 return InferSPMDExpandedLocalShape(squeeze_op); in ExpandOp()
H A Dbias_add_spmd_expander.cc105 mlir::Operation* new_local_op = InferSPMDExpandedLocalShape(op); in ExpandOp()
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/
H A Dshape_utils.h34 mlir::Operation* InferSPMDExpandedLocalShape(mlir::Operation* op);

12