Searched refs:shard_dimension (Results 1 – 4 of 4) sorted by relevance
| /external/tensorflow/tensorflow/python/tpu/ |
| D | tpu_sharding.py | 42 if self.number_of_shards is None or self.shard_dimension is None:
46 (self.number_of_shards, self.shard_dimension))
123 def shard_dimension(self): member in ShardingPolicy
127 def set_shard_dimension(self, shard_dimension): argument
142 if self._shard_dimension != shard_dimension:
145 "use %d." % (shard_dimension, self._shard_dimension))
147 self._shard_dimension = tensor_shape.as_dimension(shard_dimension)
161 if other.shard_dimension is not None:
162 self.set_shard_dimension(other.shard_dimension)
306 str(shapes), self.number_of_shards, self.shard_dimension))
|
| D | tpu_sharding_test.py | 36 self.assertEqual(p1.shard_dimension, tpu_sharding._DEFAULT_SHARD_DIMENSION)
42 self.assertEqual(p2.shard_dimension, 23)
70 self.assertEqual(p2.shard_dimension, 23)
75 self.assertEqual(p2.shard_dimension, 12)
79 self.assertEqual(p2.shard_dimension, 12)
|
| D | tpu_infeed_test.py | 54 self.assertEqual([p.shard_dimension
|
| D | tpu_feed.py | 324 return [policy.shard_dimension for policy in self._sharding_policies]
734 axis=policy.shard_dimension,
|