Home
last modified time | relevance | path

Searched defs:batch_dim (Results 1 – 25 of 34) sorted by relevance

12

/external/tensorflow/tensorflow/lite/delegates/hexagon/
Dutils_test.cc27 unsigned int batch_dim, height_dim, width_dim, depth_dim; in TEST() local
45 unsigned int batch_dim, height_dim, width_dim, depth_dim; in TEST() local
61 unsigned int batch_dim, height_dim, width_dim, depth_dim; in TEST() local
/external/tensorflow/tensorflow/core/kernels/
Dreverse_sequence_op.cc44 void CheckErrors(OpKernelContext* context, int batch_dim, int seq_dim) { in CheckErrors()
79 void CheckErrorsGPU(OpKernelContext* context, int batch_dim, int seq_dim) { in CheckErrorsGPU()
100 void CheckErrors<GPUDevice, int32>(OpKernelContext* context, int batch_dim, in CheckErrors()
106 void CheckErrors<GPUDevice, int64>(OpKernelContext* context, int batch_dim, in CheckErrors()
Dreverse_sequence_op.h32 typename TTypes<T, Dims>::ConstTensor input, int32_t batch_dim, in ReverseGenerator()
Dconv_grad_shape_utils.cc115 int batch_dim = GetTensorBatchDimIndex(num_dims, data_format); in ConvBackpropComputeDimensionsV2() local
Dscatter_nd_op.cc781 const int64_t batch_dim = (indices.dims() > 1) ? indices.dims() - 1 : 1; in ValidateUpdateShape() local
/external/tensorflow/tensorflow/compiler/xla/service/
Dspace_to_batch_converter_test.cc60 const int64_t batch_dim = reshape->operand(0) in TEST_F() local
168 const int64_t batch_dim = reshape->operand(0) in TEST_F() local
203 const int64_t batch_dim = reshape->operand(0) in TEST_F() local
Dtopk_rewriter.cc131 const int64_t batch_dim = sort_dim == 1 ? 0 : 1; in SortIsInTopK() local
197 const int64_t batch_dim = sort_dim == 1 ? 0 : 1; in TransformToCustomCall() local
Dbatch_dot_simplification.cc57 for (int64_t batch_dim : dim_numbers.lhs_batch_dimensions()) { in ElideDegenerateBatchDimensionFromBatchDot() local
Ddynamic_dimension_inference_test.cc944 WindowDimension* batch_dim = window.add_dimensions(); in TEST_F() local
994 WindowDimension* batch_dim = window.add_dimensions(); in TEST_F() local
Dspace_to_batch_converter.cc1417 const int64_t batch_dim = result.batch; in SupportedOpForPropagation() local
1585 const int64_t batch_dim = DimLookUp(permute_dims, old_batch_dim); in Propagate() local
2080 const int64_t batch_dim = DimLookUp(permute_dims, old_batch_dim); in BatchToSpace() local
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dconv_op_helpers.cc106 int batch_dim = GetTensorBatchDimIndex(num_dims, attrs.data_format); in CheckConvAttrs() local
209 int batch_dim = GetTensorBatchDimIndex(num_dims, attrs.data_format); in MakeXlaForwardConvOp() local
300 int batch_dim = GetTensorBatchDimIndex(num_dims, attrs.data_format); in MakeXlaBackpropInputConvOp() local
Dextract_image_patches_op.cc61 int batch_dim = GetTensorBatchDimIndex(num_dims, data_format); in Compile() local
Dscatter_nd_op.cc45 const int64_t batch_dim = indices_shape.dims() - 1; in ValidateUpdateShape() local
/external/tensorflow/tensorflow/core/ops/
Dimage_ops.cc30 Status SetOutputToSizedImage(InferenceContext* c, DimensionHandle batch_dim, in SetOutputToSizedImage()
228 DimensionHandle batch_dim = c->Dim(boxes, 0); in CombinedNMSShapeFn() local
822 DimensionHandle batch_dim; in __anon4ac001e21402() local
857 DimensionHandle batch_dim; in __anon4ac001e21502() local
Dboosted_trees_ops.cc495 shape_inference::DimensionHandle batch_dim = c->Dim(c->input(1), 0); in __anon16fb5b110d02() local
/external/tensorflow/tensorflow/lite/kernels/
Dreverse_sequence.cc83 int batch_dim = params->batch_dim; in ReverseSequenceImpl() local
Dreverse_sequence_test.cc33 int seq_dim, int batch_dim) { in ReverseSequenceOpModel()
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dgpu_layout_assignment.cc242 for (int64_t batch_dim : dim_nums.lhs_batch_dimensions()) { in AddBackendConstraints() local
Dgemm_thunk.cc226 for (int64_t batch_dim : batch_dimensions) { in RunGemm() local
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
Dlegalize_hlo.cc179 bool NeedsReformatTypeAndPermutation(int batch_dim, int feature_dim, in NeedsReformatTypeAndPermutation()
192 GetReformatTypeAndPermutation(int batch_dim, int feature_dim, in GetReformatTypeAndPermutation()
218 Value FormatToNHWC(Value value, int batch_dim, int feature_dim, in FormatToNHWC()
1266 const uint64_t batch_dim = 0; in IsSpatialPoolingWithoutDilation() local
/external/tensorflow/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/
Dgemm_pattern.cc232 for (auto batch_dim : dim_nums.lhs_batching_dimensions()) { in GemmOpConversionRewrite() local
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dcustom_call_handler.cc81 const int64_t batch_dim = 0; in HandleCustomCallTopK() local
/external/tensorflow/tensorflow/compiler/xla/client/lib/
Dslicing.cc286 for (int64_t batch_dim = 0; batch_dim < batch_dims; ++batch_dim) { in TorchIndexSelect() local
Dsvd_test.cc64 Array3D<float> GetUnitMatrix3D(int32_t batch_dim, int32_t mat_dim) { in GetUnitMatrix3D()
/external/tensorflow/tensorflow/cc/gradients/
Darray_grad.cc269 int batch_dim; in ReverseSequenceGrad() local

12