Home
last modified time | relevance | path

Searched refs:IsBatchParallel (Results 1 – 4 of 4) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/expansions/
H A Dconv_spmd_expander.cc58 if (input_layout.IsBatchParallel()) in VerifyConvLayout()
150 if (input_layout.IsBatchParallel()) in HandleConv()
320 if (input_layout.IsBatchParallel()) { in HandleConvBackpropFilter()
401 if (!output_layout->IsFullyReplicated() && !output_layout->IsBatchParallel()) in ExpandOp()
H A Ddataparallel_spmd_expander.cc61 if (!layouts[i].IsBatchParallel(batchable_indices.lookup(i))) return false; in AllBatchParallel()
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/cc/
H A Dtensor_layout.h302 bool IsBatchParallel() const;
304 bool IsBatchParallel(int non_batch_rank) const;
H A Dtensor_layout.cc920 bool Layout::IsBatchParallel() const { in IsBatchParallel() function in tensorflow::dtensor::Layout
935 bool Layout::IsBatchParallel(int non_batch_rank) const { in IsBatchParallel() function in tensorflow::dtensor::Layout