Searched refs:get_unsharded_shape (Results 1 – 3 of 3) sorted by relevance
| /external/tensorflow/tensorflow/python/tpu/ |
| D | tpu_sharding_test.py | 126 self.assertEqual(p.get_unsharded_shape([[4, 3], [4, 3]]), [4, 6])
128 _ = p.get_unsharded_shape([[4, 3]])
130 _ = p.get_unsharded_shape([[4, 3], [4, 3], [4, 3]])
132 _ = p.get_unsharded_shape([[4, 3], [4, 2]])
134 _ = p.get_unsharded_shape([[4, 3], "not_a_shape"])
136 _ = p.get_unsharded_shape([None, [4, 3]])
138 _ = p.get_unsharded_shape([[2], [4, 3]])
145 self.assertEqual(p.get_unsharded_shape([[]]), [])
|
| D | tpu_sharding.py | 273 def get_unsharded_shape(self, shapes): member in ShardingPolicy
|
| D | tpu_feed.py | 435 policy.get_unsharded_shape(s)
|