Home
last modified time | relevance | path

Searched refs:replica_id (Results 1 – 25 of 33) sorted by relevance

12

/external/tensorflow/tensorflow/python/distribute/
Dsummary_op_util.py45 replica_id = replica_context.replica_id_in_sync_group
46 if isinstance(replica_id, ops.Tensor):
47 replica_id = tensor_util.constant_value(replica_id)
48 return replica_id and replica_id > 0
Dmoving_averages_test.py54 replica_id = [0]
58 val = constant_op.constant([1.0 + replica_id[0], 2.0 - replica_id[0]])
59 replica_id[0] += 1
71 average_val = [1.0 + 0.5 * (replica_id[0] - 1),
72 2.0 - 0.5 * (replica_id[0] - 1)]
81 replica_id = [0]
85 val = constant_op.constant([1.0 + replica_id[0], 2.0 - replica_id[0]])
86 replica_id[0] += 1
97 average_val = [1.0 + 0.5 * (replica_id[0] - 1),
98 2.0 - 0.5 * (replica_id[0] - 1)]
Dmirrored_function_strategy_test.py48 replica_id = replica_context.replica_id_in_sync_group
49 self.assertIsInstance(replica_id, ops.Tensor)
50 self.assertIsNone(tensor_util.constant_value(replica_id))
51 return x + replica_id
Dmirrored_variable_test.py46 replica_id = ds_context.get_replica_context().replica_id_in_sync_group
47 if not isinstance(replica_id, ops.Tensor):
48 replica_id = constant_op.constant(replica_id)
49 return replica_id
198 replica_id = self.evaluate(_replica_id())
199 v = variable_scope.variable(1.0, name="foo_" + str(replica_id))
472 replica_id = self.evaluate(_replica_id())
484 v_sum.assign_add(2.0 + replica_id),
485 v_mean.assign(6.0 * replica_id)
487 all_v_sum[replica_id] = v_sum
[all …]
Dmirrored_strategy.py549 replica_id, argument
554 if replica_id == 0:
588 replica_id=i,
908 def __init__(self, dist, coord, replica_id, devices, variable_creator_fn, argument
914 self.replica_id = replica_id
960 if self.replica_id > 0:
963 self._name_scope += "replica_%d/" % self.replica_id
980 self.replica_id, dtypes.int32)), \
981 ops.device(self.devices[self.replica_id]), \
984 self._var_scope, reuse=self.replica_id > 0), \
[all …]
Dvalues.py50 replica_id = replica_context.replica_id_in_sync_group
51 if not isinstance(replica_id, int):
52 replica_id = tensor_util.constant_value(replica_id)
54 replica_id = distribute_lib.get_update_replica_id()
55 return replica_id
66 replica_id = _get_current_replica_id_as_int()
67 if replica_id is None:
70 return self._values[replica_id]
79 replica_id = _get_current_replica_id_as_int()
80 if replica_id is None:
[all …]
Dmirrored_strategy_test.py97 replica_id = int(self.evaluate(_replica_id()))
100 return list(range(replica_id))
130 replica_id = _replica_id_as_int()
131 return math_ops.cast([replica_id] * (replica_id + 1), dtype)
150 replica_id = _replica_id_as_int()
151 tensor = math_ops.cast([replica_id] * (replica_id + 1), dtype)
1384 replica_id = ds_context.get_replica_context().replica_id_in_sync_group
1385 if not isinstance(replica_id, ops.Tensor):
1386 replica_id = constant_op.constant(replica_id)
1387 return replica_id
[all …]
Dtpu_strategy.py252 for replica_id in range(device_assignment.num_replicas):
259 replica=replica_id,
418 for replica_id in range(self._num_replicas_in_sync):
419 …select_replica = lambda x: values.select_replica(replica_id, x) # pylint: disable=cell-var-from-l…
781 def replicated_fn(replica_id, replica_args, replica_kwargs): argument
783 with _TPUReplicaContext(strategy, replica_id_in_sync_group=replica_id):
866 replica_id = tensor_util.constant_value(self._replica_id_in_sync_group)
868 if replica_id is None: # Non-constant `Tensor` inside `tpu.replicate`.
872 return (ds.extended.worker_devices[replica_id],)
Dparameter_server_strategy_test.py62 replica_id = ds_context.get_replica_context().replica_id_in_sync_group
63 if isinstance(replica_id, ops.Tensor):
64 replica_id = tensor_util.constant_value(replica_id)
65 return replica_id
134 replica_id = _get_replica_id_integer()
135 last_part_device = ('device:GPU:%d' % replica_id)
279 replica_id = _get_replica_id_integer()
280 replica_compute_device = ('/device:GPU:%d' % replica_id)
287 replica_id = _get_replica_id_integer()
288 replica_variable_device = ('/device:GPU:%d' % replica_id)
Dcollective_all_reduce_strategy.py350 replica_id, argument
354 if replica_id == 0: # First replica on each worker.
391 replica_id=replica_id,
Dstrategy_test_lib.py228 replica_id = ds_context.get_replica_context().replica_id_in_sync_group
229 return summary_ops.write("a", replica_id)
261 replica_id = self.evaluate(
263 self.assertLess(replica_id, len(d.extended.worker_devices))
264 self.assertFalse(expected_devices[replica_id])
265 expected_devices[replica_id] = True
Dvalues_test.py769 replica_id = ctx.replica_id_in_sync_group
770 return v.assign(math_ops.cast(replica_id, dtypes.float32))
1228 replica_id = ctx.replica_id_in_sync_group
1229 return v.assign(math_ops.cast(replica_id, dtypes.float32))
1261 replica_id = ctx.replica_id_in_sync_group
1262 return ctx.all_reduce("SUM", v) + math_ops.cast(replica_id,
1295 replica_id = ctx.replica_id_in_sync_group
1296 return var.assign(math_ops.cast(replica_id, dtypes.float32))
Dcustom_training_loop_test.py531 for replica_id in range(distribution.num_replicas_in_sync):
534 results[0])[replica_id])
539 for replica_id in range(1, distribution.num_replicas_in_sync):
542 results[1])[replica_id])
/external/tensorflow/tensorflow/python/keras/distribute/
Dkeras_optimizer_v2_test.py65 replica_id = _replica_id()
66 return math_ops.cast(replica_id + 1, dtype=dtypes.float32) * 0.5 * var
142 replica_id = ds_context.get_replica_context().replica_id_in_sync_group
143 if not isinstance(replica_id, ops.Tensor):
144 replica_id = constant_op.constant(replica_id)
145 return replica_id
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dcollective_permute_thunk.cc91 int64 replica_id; member
163 CHECK(participants_.emplace(participant.replica_id, participant).second); in SubmitParticipant()
184 participant.replica_id, participant.stream, in SubmitParticipant()
238 TF_ASSIGN_OR_RETURN(int64 replica_id, in ExecuteOnStream()
245 if (src_dest.first == replica_id) { in ExecuteOnStream()
252 ParticipantData participant{replica_id, params.stream, src_addr, dest_addr, in ExecuteOnStream()
261 return src_dest.second == replica_id; in ExecuteOnStream()
280 participant.replica_id, participant.stream, in ExecuteOnStream()
Dreplica_id_thunk.cc30 TF_ASSIGN_OR_RETURN(int replica_id, in ExecuteOnStream()
33 params.stream->ThenMemset32(&dest_addr, replica_id, /*size=*/4); in ExecuteOnStream()
/external/tensorflow/tensorflow/compiler/xla/service/
Dcollective_ops_utils.cc60 TF_ASSIGN_OR_RETURN(int replica_id, in GetParticipatingReplicas()
66 if (absl::c_linear_search(g.replica_ids(), replica_id)) { in GetParticipatingReplicas()
68 << "Replica " << replica_id << " appears twice in replica groups"; in GetParticipatingReplicas()
73 << "Replica " << replica_id << " doesn't appear in replica groups? "; in GetParticipatingReplicas()
Dcomputation_placer.cc44 absl::optional<int> replica_id; in ReplicaIdForDeviceOrdinal() local
48 if (replica_id.has_value()) { in ReplicaIdForDeviceOrdinal()
53 replica_id = r; in ReplicaIdForDeviceOrdinal()
57 if (!replica_id.has_value()) { in ReplicaIdForDeviceOrdinal()
62 return *replica_id; in ReplicaIdForDeviceOrdinal()
Dallocation_tracker.cc186 const GlobalDataHandle& data, int replica_id) const { in ResolveForReplica()
190 if (replica_id >= replicated_buffers.size()) { in ResolveForReplica()
194 replica_id, replicated_buffers.size()); in ResolveForReplica()
196 return replicated_buffers[replica_id]; in ResolveForReplica()
Dservice.cc1033 if (arg->replica_id() < 0 || arg->replica_id() >= replica_count) { in TransferToInfeed()
1036 StrCat("The replica_id=", arg->replica_id(), in TransferToInfeed()
1045 executor = replicas[arg->replica_id()]; in TransferToInfeed()
1050 executor = replicas[arg->replica_id()]; in TransferToInfeed()
1062 if (arg->replica_id() < 0 || arg->replica_id() >= replica_count) { in TransferFromOutfeed()
1065 arg->replica_id(), replica_count); in TransferFromOutfeed()
1072 executor = replicas[arg->replica_id()]; in TransferFromOutfeed()
1077 executor = replicas[arg->replica_id()]; in TransferFromOutfeed()
Dallocation_tracker.h74 int replica_id) const;
/external/tensorflow/tensorflow/compiler/xla/client/
Dclient.h152 Status TransferToInfeed(const LiteralSlice& literal, int64 replica_id = 0,
161 const Shape* shape_with_layout, int64 replica_id = 0,
Dclient.cc96 Status Client::TransferToInfeed(const LiteralSlice& literal, int64 replica_id, in TransferToInfeed() argument
103 request.set_replica_id(replica_id); in TransferToInfeed()
119 const Shape* shape_with_layout, int64 replica_id, in TransferFromOutfeed() argument
125 request.set_replica_id(replica_id); in TransferFromOutfeed()
/external/tensorflow/tensorflow/python/tpu/
Dtpu_feed.py856 replica_id = self._device_assignment.lookup_replicas(
863 replica=replica_id, logical_core=logical_core)
867 replica=replica_id, logical_core=logical_core)
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Dcpu_runtime.cc455 xla::int32 replica_id = run_options->device_assignment() in __xla_cpu_runtime_ReplicaId() local
458 std::memcpy(output_buffer, &replica_id, 4); in __xla_cpu_runtime_ReplicaId()

12