| /external/tensorflow/tensorflow/lite/delegates/coreml/builders/ |
| D | util_test.cc | 71 std::vector<int> base_shape = {2, 2, 3, 4}; in TEST_F() local
76 SetInputShapes({base_shape, shape}); in TEST_F()
83 std::vector<int> base_shape = {1, 2, 3, 4}; in TEST_F() local
86 SetInputShapes({base_shape, shape}); in TEST_F()
93 std::vector<int> base_shape = {2, 2, 3, 4}; in TEST_F() local
95 SetInputShapes({base_shape, {1}}); in TEST_F()
101 std::vector<int> base_shape = {2, 2, 3, 4}; in TEST_F() local
105 SetInputShapes({base_shape, shape}); in TEST_F()
|
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | gather_scatter_handler.cc | 58 operand.base_shape(), operand.sharding(), partition_id, b); in IndexBoundsForGatherScatterOperandPartitionedOnTrivialSliceDims()
67 replicated_indices.base_shape().element_type(), 0, b)); in IndexBoundsForGatherScatterOperandPartitionedOnTrivialSliceDims()
69 replicated_indices.base_shape().element_type(), in IndexBoundsForGatherScatterOperandPartitionedOnTrivialSliceDims()
70 operand.base_shape().dimensions(dim), b)); in IndexBoundsForGatherScatterOperandPartitionedOnTrivialSliceDims()
75 replicated_indices.base_shape().element_type()) { in IndexBoundsForGatherScatterOperandPartitionedOnTrivialSliceDims()
77 ShapeUtil::MakeShape(replicated_indices.base_shape().element_type(), in IndexBoundsForGatherScatterOperandPartitionedOnTrivialSliceDims()
83 CreateR0WithType<int32>(replicated_indices.base_shape().element_type(), in IndexBoundsForGatherScatterOperandPartitionedOnTrivialSliceDims()
91 if (index_vector_dim < replicated_indices.base_shape().rank()) { in IndexBoundsForGatherScatterOperandPartitionedOnTrivialSliceDims()
112 replicated_indices.base_shape(), min_indices[0], {index_vector_dim})); in IndexBoundsForGatherScatterOperandPartitionedOnTrivialSliceDims()
114 replicated_indices.base_shape(), max_indices[0], {index_vector_dim})); in IndexBoundsForGatherScatterOperandPartitionedOnTrivialSliceDims()
[all …]
|
| D | convolution_handler.cc | 58 lhs.base_shape().dimensions(dnums.input_batch_dimension()); in PartitionConvolutionWithBatchGroupCount()
60 rhs.base_shape().dimensions(dnums.kernel_output_feature_dimension()); in PartitionConvolutionWithBatchGroupCount()
84 std::vector<int64> lhs_to_output_indices(lhs.base_shape().rank(), -1); in PartitionConvolutionWithBatchGroupCount()
149 lhs.base_shape().dimensions(dnums.input_feature_dimension()); in PartitionConvolutionWithFeatureGroupCount()
151 rhs.base_shape().dimensions(dnums.kernel_output_feature_dimension()); in PartitionConvolutionWithFeatureGroupCount()
270 if (ShapeSizeInBytes(lhs.base_shape()) < ShapeSizeInBytes(rhs.base_shape())) { in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
332 CeilOfRatio(lhs.base_shape().dimensions(lhs_dimension), shard_count); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
334 CeilOfRatio(rhs.base_shape().dimensions(rhs_dimension), shard_count); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
375 int64 base = lhs.base_shape().dimensions(lhs_dimension); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
486 rhs_with_halo, rhs.base_shape(), left_halo_size_functions[dim], in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
[all …]
|
| D | spmd_partitioner_util.cc | 276 Shape GetPaddedShapeForUnevenPartitioning(const Shape& base_shape, in GetPaddedShapeForUnevenPartitioning() argument
279 return base_shape; in GetPaddedShapeForUnevenPartitioning()
281 if (EvenlyPartitions(base_shape, sharding)) { in GetPaddedShapeForUnevenPartitioning()
282 return base_shape; in GetPaddedShapeForUnevenPartitioning()
284 auto shard_shape = MakePartitionedShape(base_shape, sharding); in GetPaddedShapeForUnevenPartitioning()
285 Shape padded_base_shape = base_shape; in GetPaddedShapeForUnevenPartitioning()
409 HloInstruction* hlo, const Shape& base_shape, in TileToPartialReplicateHaloExchange() argument
416 GetPaddedShapeForUnevenPartitioning(base_shape, src_sharding); in TileToPartialReplicateHaloExchange()
419 GetPaddedShapeForUnevenPartitioning(base_shape, dst_sharding); in TileToPartialReplicateHaloExchange()
503 HloInstruction* hlo, const Shape& base_shape, in PadFromPartialReplicateShape() argument
[all …]
|
| D | dot_handler.cc | 492 dims_mapping, lhs.base_shape().rank(), rhs.base_shape().rank(), in PartitionBaseCase()
1474 ShapeSizeInBytes(rhs.base_shape()), in PartitionBaseCase()
1475 ShapeSizeInBytes(lhs.base_shape()), in PartitionBaseCase()
1501 if (ShapeSizeInBytes(lhs.base_shape()) < in PartitionBaseCase()
1502 ShapeSizeInBytes(rhs.base_shape())) { in PartitionBaseCase()
1514 lhs_contracting_dims.reserve(lhs.base_shape().rank()); in PartitionBaseCase()
1594 ShapeUtil::ElementsIn(rhs.base_shape()) > in PartitionBaseCase()
1599 ShapeUtil::ElementsIn(lhs.base_shape()) > in PartitionBaseCase()
1624 lhs_contracting_dims.reserve(lhs.base_shape().rank()); in PartitionBaseCase()
1674 ? std::vector<int64>(lhs.base_shape().rank(), 1) in PartitionDotGroupOnBatch()
[all …]
|
| D | spmd_partitioner_util.h | 121 Shape GetPaddedShapeForUnevenPartitioning(const Shape& base_shape,
253 HloInstruction* hlo, const Shape& base_shape,
367 HloInstruction* hlo, const Shape& base_shape,
390 HloInstruction* hlo, const Shape& base_shape,
|
| D | spmd_partitioner.h | 293 PartitionedHlo(HloInstruction* hlo, Shape base_shape, PartitioningState state) in PartitionedHlo() argument
294 : hlo_(hlo), base_shape_(base_shape), state_(std::move(state)) { in PartitionedHlo()
324 const Shape& base_shape() const { return base_shape_; } in base_shape() function
|
| D | spmd_partitioner.cc | 464 return PartitionedHlo(partially_sharded, base_shape(), state_); in ReshardNoCache()
874 target_shape.set_dimensions(i, base_shape().dimensions(i)); in ReplicatePartial()
961 auto base_shape = MakePartitionedShape(base_shape_, target); in ReshardToPartialReplicateWithAllGather() local
966 auto partial_replicate_hlo = PartitionedHlo(halo_exchange_hlo, base_shape, in ReshardToPartialReplicateWithAllGather()
1597 operand.base_shape().dimensions(i)); in HandleSlice()
1862 value_gte, partitioned_topk.base_shape().tuple_shapes(0), in HandleCustomCall()
1889 index_gte, partitioned_topk.base_shape().tuple_shapes(1), in HandleCustomCall()
2021 input_major_dims_size *= operand.base_shape().dimensions(i); in HandleReshape()
2043 int64 input_dim_size = operand.base_shape().dimensions(input_sharded_dim); in HandleReshape()
2046 MakePartitionedShape(operand.base_shape(), operand.sharding()); in HandleReshape()
[all …]
|
| /external/tensorflow/tensorflow/lite/testing/nnapi_tflite_zip_tests/ |
| D | test_manifest.txt | 65 concat/concat_axis=0,base_shape=[1,3,4,3],num_tensors=1,type=tf.float32
66 concat/concat_axis=0,base_shape=[1,3,4,3],num_tensors=1,type=tf.uint8
67 DISABLED_concat/concat_axis=0,base_shape=[1,3,4,3],num_tensors=1,type=tf.int32
68 DISABLED_concat/concat_axis=0,base_shape=[1,3,4,3],num_tensors=1,type=tf.int64
69 concat/concat_axis=0,base_shape=[3,4],num_tensors=1,type=tf.float32
70 concat/concat_axis=0,base_shape=[3,4],num_tensors=1,type=tf.uint8
71 DISABLED_concat/concat_axis=0,base_shape=[3,4],num_tensors=1,type=tf.int32
72 DISABLED_concat/concat_axis=0,base_shape=[3,4],num_tensors=1,type=tf.int64
73 concat/concat_axis=0,base_shape=[1,3,4,3],num_tensors=2,type=tf.float32
74 concat/concat_axis=0,base_shape=[1,3,4,3],num_tensors=2,type=tf.uint8
[all …]
|
| /external/tensorflow/tensorflow/lite/kernels/ |
| D | mul_test.cc | 168 const std::vector<int> base_shape = {2, 3, 1, 2}; in TEST() local
186 {TensorType_FLOAT32, base_shape}, {TensorType_FLOAT32, test_shapes[i]}, in TEST()
202 {TensorType_FLOAT32, test_shapes[i]}, {TensorType_FLOAT32, base_shape}, in TEST()
424 const std::vector<int> base_shape = {2, 3, 1, 2}; in QuantizedWithMixedBroadcast() local
441 QuantizedMulOpModel model_fixture({tensor_type, base_shape, -3.f, 3.f}, in QuantizedWithMixedBroadcast()
459 {tensor_type, base_shape, -3.f, 3.f}, in QuantizedWithMixedBroadcast()
|
| D | add_test.cc | 182 const std::vector<int> base_shape = {2, 3, 1, 2}; in TEST() local
200 {TensorType_FLOAT32, base_shape}, {TensorType_FLOAT32, test_shapes[i]}, in TEST()
215 {TensorType_FLOAT32, test_shapes[i]}, {TensorType_FLOAT32, base_shape}, in TEST()
471 const std::vector<int> base_shape = {2, 3, 1, 2}; in QuantizedWithMixedBroadcast() local
488 QuantizedAddOpModel model_fixture({tensor_type, base_shape, -3.f, 3.f}, in QuantizedWithMixedBroadcast()
506 {tensor_type, base_shape, -3.f, 3.f}, in QuantizedWithMixedBroadcast()
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | shape_inference.cc | 161 StatusOr<Shape> InferWindowOutputShape(const Shape& base_shape, in InferWindowOutputShape() argument
165 if (window.dimensions_size() != base_shape.rank()) { in InferWindowOutputShape()
168 window.dimensions_size(), base_shape.rank()); in InferWindowOutputShape()
203 ShapeUtil::GetDimension(base_shape, i), dim.base_dilation()); in InferWindowOutputShape()
211 output_is_dynamic[i] = base_shape.is_dynamic_dimension(i); in InferWindowOutputShape()
1813 Shape base_shape = in InferConvolveShape() local
1817 InferWindowOutputShape(base_shape, window, lhs.element_type(), in InferConvolveShape()
|
| D | hlo_evaluator_typed_visitor.h | 2840 const Shape& window_shape, const Window& window, const Shape& base_shape,
2843 const int64 rank = base_shape.rank();
2872 if (base_index[i] < 0 || base_index[i] >= base_shape.dimensions(i)) {
|