Searched refs:partition_shape (Results 1 – 6 of 6) sorted by relevance
| /external/tensorflow/tensorflow/cc/gradients/ |
| D | data_flow_grad_test.cc | 54 TensorShape partition_shape({2, 2}); in TEST_F() local
56 {partition_shape, partition_shape, partition_shape}); in TEST_F()
|
| /external/tensorflow/tensorflow/python/ops/ |
| D | init_ops_v2_test.py | 80 partition_shape = (2, 2)
87 partition_shape=partition_shape,
89 self.assertEqual(got.shape, partition_shape)
91 got, array_ops.slice(full_value, partition_offset, partition_shape))
138 init((4, 2), dtype=dtypes.float32, partition_shape=(2, 2))
375 partition_shape = (100, 100)
384 x = self.evaluate(init(shape, partition_shape=partition_shape))
387 self.assertEqual(x.shape, partition_shape)
462 init((4, 2), dtype=dtypes.float32, partition_shape=(2, 2))
524 init((4, 2), dtype=dtypes.float32, partition_shape=(2, 2))
|
| /external/tensorflow/tensorflow/core/ops/ |
| D | ragged_conversion_ops.cc | 72 TensorShapeProto partition_shape; in ValidateRowPartitionTypesAndShapes() local
73 c->ShapeHandleToProto(c->input(3 + i), &partition_shape); in ValidateRowPartitionTypesAndShapes()
74 if (partition_shape.unknown_rank()) { in ValidateRowPartitionTypesAndShapes()
78 if (partition_shape.dim_size() != 0) { in ValidateRowPartitionTypesAndShapes()
82 if (partition_shape.dim_size() != 1) { in ValidateRowPartitionTypesAndShapes()
|
| /external/tensorflow/tensorflow/python/distribute/ |
| D | parameter_server_strategy_v2.py | 628 partition_shape = (offsets[shard_index + 1] -
636 partition_shape=partition_shape,
644 if value.shape == partition_shape:
659 shape=tensor_shape.as_shape(partition_shape),
|
| /external/tensorflow/tensorflow/python/keras/ |
| D | initializers_test.py | 270 shape=(4, 2), partition_shape=(2, 2), partition_offset=(0, 0))
282 shape=(4, 2), partition_shape=(2, 2), partition_offset=(0, 0))
|
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner_util.cc | 189 auto partition_shape = shape; in MakeNonPaddedShapeForGivenPartition() local
196 partition_shape.set_dimensions(i, tile_limit[i] - tile_offset[i]); in MakeNonPaddedShapeForGivenPartition()
198 partition_shape.set_dimensions(i, 0); in MakeNonPaddedShapeForGivenPartition()
201 return partition_shape; in MakeNonPaddedShapeForGivenPartition()
|