Home
last modified time | relevance | path

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

12

/aosp_15_r20/external/tensorflow/tensorflow/core/kernels/
H A 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_t>(OpKernelContext* context, int batch_dim, in CheckErrors()
H A Dscatter_nd_util.cc29 const int64_t batch_dim = in ValidateScatterNdUpdateShape() local
H A Dreverse_sequence_op.h32 typename TTypes<T, Dims>::ConstTensor input, int32_t batch_dim, in ReverseGenerator()
H A Dconv_grad_shape_utils.cc115 int batch_dim = GetTensorBatchDimIndex(num_dims, data_format); in ConvBackpropComputeDimensionsV2() local
/aosp_15_r20/external/tensorflow/tensorflow/lite/delegates/hexagon/
H A 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
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dspace_to_batch_converter_test.cc62 const int64_t batch_dim = previous_reshape->operand(0) in TEST_F() local
205 const int64_t batch_dim = previous_reshape->operand(0) in TEST_F() local
242 const int64_t batch_dim = previous_reshape->operand(0) in TEST_F() local
H A Dtopk_rewriter.cc139 const int64_t batch_dim = sort_dim == 1 ? 0 : 1; in SortIsInTopK() local
207 const int64_t batch_dim = sort_dim == 1 ? 0 : 1; in TransformToCustomCall() local
H A Dbatch_dot_simplification.cc56 for (int64_t batch_dim : dim_numbers.lhs_batch_dimensions()) { in ElideDegenerateBatchDimensionFromBatchDot() local
H A Ddynamic_dimension_inference_test.cc951 WindowDimension* batch_dim = window.add_dimensions(); in TEST_F() local
1001 WindowDimension* batch_dim = window.add_dimensions(); in TEST_F() local
/aosp_15_r20/external/tensorflow/tensorflow/core/ops/
H A 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 __anon269972ec1402() local
857 DimensionHandle batch_dim; in __anon269972ec1502() local
/aosp_15_r20/external/pytorch/torch/csrc/api/include/torch/nn/utils/
H A Drnn.h211 int64_t batch_dim = batch_first ? 0 : 1; variable
272 int64_t batch_dim = batch_first ? 0 : 1; variable
/aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
H A 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
H A Dextract_image_patches_op.cc61 int batch_dim = GetTensorBatchDimIndex(num_dims, data_format); in Compile() local
H A Dscatter_nd_op.cc46 const int64_t batch_dim = indices_shape.dims() - 1; in ValidateUpdateShape() local
/aosp_15_r20/external/tensorflow/tensorflow/lite/kernels/
H A Dreverse_sequence.cc83 int batch_dim = params->batch_dim; in ReverseSequenceImpl() local
H A Dreverse_sequence_test.cc33 int seq_dim, int batch_dim) { in ReverseSequenceOpModel()
/aosp_15_r20/external/pytorch/aten/src/ATen/native/
H A Dim2col_shape_check.h55 int64_t batch_dim = (ndim == 3) ? 0 : -1; in col2im_shape_check() local
H A DLegacyBatching.cpp14 Tensor _add_batch_dim(const Tensor& self, int64_t batch_dim, int64_t level) { in _add_batch_dim()
/aosp_15_r20/external/pytorch/aten/src/ATen/
H A DSparseCsrTensorImpl.h69 inline int64_t batch_dim() const noexcept { in batch_dim() function
H A DSparseCsrTensorImpl.cpp115 auto batch_dim = size.size() - sparse_dim - dense_dim; in resize_and_clear_() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
H A Dlegalize_hlo.cc500 bool NeedsReformatTypeAndPermutation(int batch_dim, int feature_dim, in NeedsReformatTypeAndPermutation()
513 GetReformatTypeAndPermutation(int batch_dim, int feature_dim, in GetReformatTypeAndPermutation()
539 Value FormatToNHWC(Value value, int batch_dim, int feature_dim, in FormatToNHWC()
2318 const uint64_t batch_dim = 0; in IsSpatialPoolingWithoutDilation() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/
H A Dcustom_call_handler.cc82 const int64_t batch_dim = 0; in HandleCustomCallTopK() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/
H A Dgemm_pattern.cc232 for (auto batch_dim : dim_nums.lhs_batching_dimensions()) { in GemmOpConversionRewrite() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/lib/
H A Dsvd_test.cc64 Array3D<float> GetUnitMatrix3D(int32_t batch_dim, int32_t mat_dim) { in GetUnitMatrix3D()
H A Dslicing.cc289 for (int64_t batch_dim = 0; batch_dim < batch_dims; ++batch_dim) { in TorchIndexSelect() local

12