Home
last modified time | relevance | path

Searched refs:offset_shape (Results 1 – 6 of 6) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Dbatchnorm_expander_test.cc59 Shape offset_shape = ShapeUtil::MakeShape(F32, {2}); in TEST_F() local
69 HloInstruction::CreateParameter(2, offset_shape, "offset")); in TEST_F()
72 ShapeUtil::MakeTupleShape({input_shape, scale_shape, offset_shape}), in TEST_F()
Dshape_inference.cc1212 const Shape& offset_shape, int64 feature_index) { in InferBatchNormTrainingShape() argument
1216 ExpectArray(offset_shape, "offset input of batch norm training")); in InferBatchNormTrainingShape()
1222 TF_RET_CHECK(ShapeUtil::ValidateShapeWithOptionalLayout(offset_shape) == in InferBatchNormTrainingShape()
1249 if (offset_shape.rank() != 1) { in InferBatchNormTrainingShape()
1253 offset_shape.rank()); in InferBatchNormTrainingShape()
1270 if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(offset_shape, in InferBatchNormTrainingShape()
1276 PrimitiveType_Name(offset_shape.element_type()), in InferBatchNormTrainingShape()
1294 if (ShapeUtil::GetDimension(offset_shape, 0) != feature_count) { in InferBatchNormTrainingShape()
1299 ShapeUtil::GetDimension(offset_shape, 0), feature_count); in InferBatchNormTrainingShape()
1317 const Shape& offset_shape, const Shape& mean_shape, in InferBatchNormInferenceShape() argument
[all …]
Dshape_inference.h90 const Shape& offset_shape,
97 const Shape& offset_shape, const Shape& mean_shape,
/external/tensorflow/tensorflow/lite/toco/graph_transformations/
Dresolve_batch_normalization.cc97 const auto& offset_shape = offset_array.shape(); in Run() local
105 CHECK(mean_shape.dims() == offset_shape.dims()); in Run()
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner.cc2399 const Shape& offset_shape = per_partition_size_hlo->shape(); in HandleDynamicUpdateSlice() local
2401 offset_shape, HloOpcode::kMultiply, partition_ordinals[dim], in HandleDynamicUpdateSlice()
2411 offset_shape, HloOpcode::kMultiply, in HandleDynamicUpdateSlice()
2413 offset_shape, HloOpcode::kAdd, partition_ordinals[dim], in HandleDynamicUpdateSlice()
/external/tensorflow/tensorflow/compiler/xla/client/
Dxla_builder.cc2648 TF_ASSIGN_OR_RETURN(const Shape* offset_shape, GetShapePtr(offset)); in BatchNormTraining()
2652 *operand_shape, *scale_shape, *offset_shape, feature_index)); in BatchNormTraining()
2671 TF_ASSIGN_OR_RETURN(const Shape* offset_shape, GetShapePtr(offset)); in BatchNormInference()
2676 *operand_shape, *scale_shape, *offset_shape, in BatchNormInference()