Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/python/tpu/
Dtpu_sharding.py41 if self.number_of_shards is None or self.shard_dimension is None:
45 (self.number_of_shards, self.shard_dimension))
65 def number_of_shards(self): member in ShardingPolicy
69 def set_number_of_shards(self, number_of_shards): argument
83 if self._number_of_shards != number_of_shards:
86 "frozen to use %d." % (number_of_shards, self._number_of_shards))
88 if number_of_shards > 0:
89 self._number_of_shards = number_of_shards
93 str(number_of_shards))
132 if other.number_of_shards is not None:
[all …]
Dtpu_function.py33 def number_of_shards(self): member in TpuContext
36 def set_number_of_shards(self, number_of_shards): argument
37 self._number_of_shards = number_of_shards
46 def tpu_shard_context(number_of_shards): argument
47 if _current_tpu_context.number_of_shards is not None:
50 _current_tpu_context.set_number_of_shards(number_of_shards)
Dtpu_infeed_test.py37 self.assertEqual(i.number_of_shards, None)
44 self.assertEqual(i.number_of_shards, None)
49 self.assertEqual(i.number_of_shards, None)
84 self.assertEqual(i.number_of_shards, 2)
86 self.assertEqual(i.number_of_shards, 3)
93 self.assertEqual(i.number_of_shards, 2)
107 self.assertEqual(i.number_of_shards, 2)
124 self.assertEqual(i.number_of_shards, 2)
Dtpu_feed.py335 def number_of_shards(self): member in InfeedQueue
342 return self._sharding_policies[0].number_of_shards
344 def set_number_of_shards(self, number_of_shards): argument
356 policy.set_number_of_shards(number_of_shards)
403 number_of_shards = len(input_tensors)
404 self.set_number_of_shards(number_of_shards)
421 for i in xrange(1, self.number_of_shards):
606 for (shard, index) in zip(sharded_inputs, xrange(self.number_of_shards))
697 if self.number_of_shards == 1:
708 self.number_of_shards,
[all …]
Dtraining_loop.py111 number_of_shards = tpu_function.get_tpu_context().number_of_shards
112 if number_of_shards is None:
117 infeed_queue.set_number_of_shards(number_of_shards)
164 num_replicas = tpu_function.get_tpu_context().number_of_shards
Dtpu_sharding_test.py34 self.assertEqual(p1.number_of_shards,
41 self.assertEqual(p2.number_of_shards, 17)
69 self.assertEqual(p2.number_of_shards, 17)
74 self.assertEqual(p2.number_of_shards, 17)
78 self.assertEqual(p2.number_of_shards, 17)
Dtpu_optimizer.py123 num_shards = tpu_function.get_tpu_context().number_of_shards
Dfeature_column.py435 return tpu_function.get_tpu_context().number_of_shards is None
/external/tensorflow/tensorflow/python/tpu/ops/
Dtpu_ops.py34 num_shards = tpu_function.get_tpu_context().number_of_shards
/external/tensorflow/tensorflow/contrib/tpu/python/tpu/
Dkeras_support.py350 num_shards = tpu_function.get_tpu_context().number_of_shards