Home
last modified time | relevance | path

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

/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/expansions/
H A Dconv_spmd_expander.cc168 const auto input_num_shards = input_layout.num_shards(); in HandleConv() local
223 int input_size = input_local_size * input_num_shards[curr_input_dim]; in HandleConv()
257 halo_size / (input_num_shards[curr_input_dim] - 1); in HandleConv()