Home
last modified time | relevance | path

Searched refs:GetTensorBatchDimIndex (Results 1 – 10 of 10) sorted by relevance

/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()
209 int batch_dim = GetTensorBatchDimIndex(num_dims, attrs.data_format); in MakeXlaForwardConvOp()
300 int batch_dim = GetTensorBatchDimIndex(num_dims, attrs.data_format); in MakeXlaBackpropInputConvOp()
430 int n_dim = GetTensorBatchDimIndex(num_dims, attrs.data_format); in MakeXlaBackpropFilterConvOp()
H A Dextract_image_patches_op.cc61 int batch_dim = GetTensorBatchDimIndex(num_dims, data_format); in Compile()
H A Dpooling_ops.cc148 int batch_dimension = GetTensorBatchDimIndex(num_dims, data_format); in XlaTensorFormat()
/aosp_15_r20/external/tensorflow/tensorflow/core/util/
H A Dtensor_format.h173 inline int GetTensorBatchDimIndex(int num_dims, TensorFormat format) { in GetTensorBatchDimIndex() function
527 dim_sizes[GetTensorBatchDimIndex(dims, format)] = N; in ShapeFromFormat()
H A Dpadding.cc52 const int32_t batch_index = GetTensorBatchDimIndex(num_dims, data_format); in CheckValidPadding()
H A Dtensor_format_test.cc178 EXPECT_EQ(GetTensorBatchDimIndex(num_dims, format), tdm.n()); in RunDimensionIndexesTest()
/aosp_15_r20/external/tensorflow/tensorflow/core/kernels/
H A Dconv_grad_shape_utils.cc115 int batch_dim = GetTensorBatchDimIndex(num_dims, data_format); in ConvBackpropComputeDimensionsV2()
/aosp_15_r20/external/tensorflow/tensorflow/core/framework/
H A Dcommon_shape_fns.cc503 dims_actual[GetTensorBatchDimIndex(num_dims, format)] = context->MakeDim(N); in MakeShapeFromFormat()
530 *batch_dim = context->Dim(shape, GetTensorBatchDimIndex(rank, format)); in DimensionsFromShape()
559 out_dims[tensorflow::GetTensorBatchDimIndex(rank, format)] = batch_dim; in ShapeFromDimensions()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/ir/
H A Dtf_ops_a_m.cc1761 return_shape[GetTensorBatchDimIndex(num_dims, format)] = in inferConvReturnTypeComponents()
1762 input_ty.getDimSize(GetTensorBatchDimIndex(num_dims, format)); in inferConvReturnTypeComponents()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/
H A Dlegalize_tf.cc990 int64_t batch_dim = GetTensorBatchDimIndex(num_dims, format); in GetConvDimensionNumbersAttr()
4867 int batch_dim = GetTensorBatchDimIndex(num_dims, data_format); in matchAndRewrite()
5248 tensorflow::GetTensorBatchDimIndex(num_dims, data_format); in matchAndRewrite()