Searched refs:get_sharded_shape (Results 1 – 3 of 3) sorted by relevance
| /external/tensorflow/tensorflow/python/tpu/ |
| D | tpu_sharding_test.py | 95 self.assertEqual(p.get_sharded_shape([4, 9]), [4, 3])
100 _ = p.get_sharded_shape([4, 9], shard_index=4)
102 _ = p.get_sharded_shape([4, 9], shard_index=-1)
104 _ = p.get_sharded_shape("not_a_shape")
106 _ = p.get_sharded_shape(tensor_shape.TensorShape(None))
108 _ = p.get_sharded_shape([4, 10], shard_index=-1)
144 self.assertEqual(p.get_sharded_shape([]), [])
|
| D | tpu_feed.py | 221 _ = policy.get_sharded_shape(shape)
501 policy.get_unpartitioned_shape(policy.get_sharded_shape(shape))
808 policy.get_sharded_shape(shape)
|
| D | tpu_sharding.py | 192 def get_sharded_shape(self, shape, shard_index=None): member in ShardingPolicy
|