Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/python/tpu/
Dtpu_sharding.py42 if self.number_of_shards is None or self.shard_dimension is None:
46 (self.number_of_shards, self.shard_dimension))
66 def number_of_shards(self): member in ShardingPolicy
70 def set_number_of_shards(self, number_of_shards): argument
84 if self._number_of_shards != number_of_shards:
87 "frozen to use %d." % (number_of_shards, self._number_of_shards))
89 if number_of_shards > 0:
90 self._number_of_shards = number_of_shards
94 str(number_of_shards))
159 if other.number_of_shards is not None:
[all …]
Dtpu_function.py34 def number_of_shards(self): member in TpuContext
37 def set_number_of_shards(self, number_of_shards): argument
38 self._number_of_shards = number_of_shards
47 def tpu_shard_context(number_of_shards): argument
49 if _current_tpu_context.number_of_shards is not None:
57 _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.py352 def number_of_shards(self): member in InfeedQueue
359 return self._sharding_policies[0].number_of_shards
361 def set_number_of_shards(self, number_of_shards): argument
373 policy.set_number_of_shards(number_of_shards)
421 number_of_shards = len(input_tensors)
422 self.set_number_of_shards(number_of_shards)
439 for i in xrange(1, self.number_of_shards):
631 for (shard, index) in zip(sharded_inputs, xrange(self.number_of_shards))
722 if self.number_of_shards == 1:
733 self.number_of_shards,
[all …]
Dtraining_loop.py119 number_of_shards = tpu_function.get_tpu_context().number_of_shards
120 if number_of_shards is None:
125 infeed_queue.set_number_of_shards(number_of_shards)
172 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.py144 num_shards = tpu_function.get_tpu_context().number_of_shards
Dfeature_column.py675 return tpu_function.get_tpu_context().number_of_shards is None
/external/tensorflow/tensorflow/python/tpu/ops/
Dtpu_ops.py35 num_shards = tpu_function.get_tpu_context().number_of_shards