Searched refs:batch_dimension (Results 1 – 3 of 3) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/client/lib/ |
| D | pooling.h | 27 TensorFormat(int batch_dimension, int feature_dimension, in TensorFormat() argument
29 : batch_dimension_(batch_dimension), in TensorFormat()
34 int batch_dimension() const { return batch_dimension_; } in batch_dimension() function
|
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | pooling_ops.cc | 140 int batch_dimension = GetTensorBatchDimIndex(num_dims, data_format); in XlaTensorFormat() local
147 return xla::TensorFormat(/*batch_dimension=*/batch_dimension, in XlaTensorFormat()
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | space_to_batch_converter.cc | 169 HloInstruction* activations, int64 batch_dimension, int64 old_batch_size,
570 HloInstruction* activations, int64 batch_dimension, int64 old_batch_size, in IncreaseSpatialSizeOnSpaceToBatchedShape() argument
572 CHECK_EQ(batch_dimension + 1, spatial_dimension); in IncreaseSpatialSizeOnSpaceToBatchedShape()
576 const int64 new_batch_size = activations->shape().dimensions(batch_dimension); in IncreaseSpatialSizeOnSpaceToBatchedShape()
584 new_dimensions[batch_dimension] = old_batch_size; in IncreaseSpatialSizeOnSpaceToBatchedShape()
611 reshape_back_dims[batch_dimension] = new_batch_size; in IncreaseSpatialSizeOnSpaceToBatchedShape()
|