Searched refs:num_per_replica_inputs (Results 1 – 2 of 2) sorted by relevance
787 int num_per_replica_inputs = in RemoveUnusedXlaInput() local815 if (index < num_per_replica_inputs) { in RemoveUnusedXlaInput()817 } else if (index < num_per_replica_inputs + num_distributed_vars) { in RemoveUnusedXlaInput()853 num_per_replica_inputs - num_removed_per_replica_inputs; in RemoveUnusedXlaInput()855 for (int j = 0; j < num_per_replica_inputs; j++) { in RemoveUnusedXlaInput()858 const Edge* e = input_edges.at(i * num_per_replica_inputs + j); in RemoveUnusedXlaInput()866 const Edge* e = input_edges.at(i * num_per_replica_inputs + j); in RemoveUnusedXlaInput()873 for (int i = num_replicas * num_per_replica_inputs; in RemoveUnusedXlaInput()876 num_per_replica_inputs + i - num_replicas * num_per_replica_inputs; in RemoveUnusedXlaInput()896 for (int j = 0; j < num_per_replica_inputs; j++) { in RemoveUnusedXlaInput()[all …]
3988 const int num_per_replica_inputs = input_types.size() - num_distributed_vars; in GetIOTypes() local3990 if (num_per_replica_inputs % num_replicas != 0) { in GetIOTypes()3992 "Number of inputs to TPUReplicate (", num_per_replica_inputs, in GetIOTypes()4012 num_replicas, num_per_replica_inputs / num_replicas, num_distributed_vars, in GetIOTypes()