Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/core/kernels/sparse/
Dsparse_mat_mul_op.cc138 const int row_dim = (rank == 2) ? 0 : 1; in Compute() local
140 SwapDimSizes(row_dim, row_dim + 1, &a_shape); in Compute()
142 SwapDimSizes(row_dim, row_dim + 1, &b_shape); in Compute()
145 ctx, a_shape.dim_size(row_dim + 1) == b_shape.dim_size(row_dim), in Compute()
157 output_shape_vec(row_dim) = a_shape.dim_size(row_dim); in Compute()
158 output_shape_vec(row_dim + 1) = b_shape.dim_size(row_dim + 1); in Compute()
174 const int64 num_output_rows = output_shape_vec(row_dim); in Compute()
177 static_cast<double>(a_shape.dim_size(row_dim) * batch_size); in Compute()
180 static_cast<double>(b_shape.dim_size(row_dim) * batch_size); in Compute()
343 const int row_dim = (rank == 2) ? 0 : 1; in Compute() local
[all …]
Dadd_op.cc280 const int row_dim = a.dense_shape_host.size() == 2 ? 0 : 1; in GetOutputStructure() local
281 DCHECK_EQ(m, a.dense_shape_host(row_dim)); in GetOutputStructure()
282 DCHECK_EQ(m, b.dense_shape_host(row_dim)); in GetOutputStructure()
287 const int n = a.dense_shape_host(row_dim + 1); in GetOutputStructure()
288 DCHECK_EQ(n, b.dense_shape_host(row_dim + 1)); in GetOutputStructure()
308 const int row_dim = a.dense_shape_host.size() == 2 ? 0 : 1; in Compute() local
309 DCHECK_EQ(m, a.dense_shape_host(row_dim)); in Compute()
310 DCHECK_EQ(m, b.dense_shape_host(row_dim)); in Compute()
314 const int n = a.dense_shape_host(row_dim + 1); in Compute()
315 DCHECK_EQ(n, b.dense_shape_host(row_dim + 1)); in Compute()
Dmat_mul_op.cc515 const int row_dim = (rank == 2) ? 0 : 1; in Compute() local
517 this->transpose_a_ ? row_dim + 1 : row_dim); in Compute()
519 b_t.shape().dim_size(this->transpose_b_ ? row_dim + 1 : row_dim); in Compute()
521 b_t.dim_size(this->transpose_b_ ? row_dim : row_dim + 1); in Compute()
534 const int64 c_matrix_lhs = c_shape.dim_size(row_dim); in Compute()
535 const int64 c_matrix_rhs = c_shape.dim_size(row_dim + 1); in Compute()
544 TTypes<int64>::ConstVec a_dense_shape_comp(a_dense_shape.data() + row_dim, in Compute()
638 b_t_transposed_shape.AddDim(b_t.dim_size(row_dim + 1)); in Compute()
639 b_t_transposed_shape.AddDim(b_t.dim_size(row_dim)); in Compute()
654 a_input_dense_shape.data() + row_dim, 2); in Compute()
Dsparse_cholesky_op.cc244 const int row_dim = (rank == 2) ? 0 : 1; in ValidateInputs() local
246 *num_rows = dense_shape_vec(row_dim); in ValidateInputs()
247 const int64 num_cols = dense_shape_vec(row_dim + 1); in ValidateInputs()
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dgemm_thunk.cc181 int64 row_dim = dim_nums.lhs_batch_dimensions_size(); in RunGemm() local
188 CHECK_NE(row_dim, batch_dim); in RunGemm()
195 CHECK_LT(shape->layout().minor_to_major(row_dim), 2); in RunGemm()
203 int64 output_num_rows = output_shape.dimensions(row_dim); in RunGemm()
228 bool is_row_major = LayoutUtil::Minor(shape.layout(), row_dim) != 0; in RunGemm()
229 bool layout_mismatch = LayoutUtil::Minor(shape.layout(), row_dim) != in RunGemm()
230 LayoutUtil::Minor(output_shape.layout(), row_dim); in RunGemm()
233 shape.dimensions(row_dim + static_cast<int64>(is_row_major)), in RunGemm()
234 shape.dimensions(row_dim + static_cast<int64>(!is_row_major))}; in RunGemm()
238 lhs_buffer, lhs_shape, dim_nums.lhs_contracting_dimensions(0) == row_dim); in RunGemm()
[all …]
/external/tensorflow/tensorflow/core/kernels/data/experimental/
Ddense_to_sparse_batch_dataset_op.cc251 for (int64_t row_dim = strides.size() - 2; row_dim >= 0; in GetNextInternal() local
252 --row_dim) { in GetNextInternal()
253 strides[row_dim] = in GetNextInternal()
254 strides[row_dim + 1] * t.shape().dim_size(row_dim + 1); in GetNextInternal()