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.cc298 const int row_dim = a.dense_shape_host.size() == 2 ? 0 : 1; in GetWorkspaceSize() local
299 DCHECK_EQ(m, a.dense_shape_host(row_dim)); in GetWorkspaceSize()
300 DCHECK_EQ(m, b.dense_shape_host(row_dim)); in GetWorkspaceSize()
304 const int n = a.dense_shape_host(row_dim + 1); in GetWorkspaceSize()
305 DCHECK_EQ(n, b.dense_shape_host(row_dim + 1)); in GetWorkspaceSize()
326 const int row_dim = a.dense_shape_host.size() == 2 ? 0 : 1; in GetOutputStructure() local
327 DCHECK_EQ(m, a.dense_shape_host(row_dim)); in GetOutputStructure()
328 DCHECK_EQ(m, b.dense_shape_host(row_dim)); in GetOutputStructure()
333 const int n = a.dense_shape_host(row_dim + 1); in GetOutputStructure()
334 DCHECK_EQ(n, b.dense_shape_host(row_dim + 1)); in GetOutputStructure()
[all …]
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()
549 TTypes<int64>::ConstVec a_dense_shape_comp(a_dense_shape.data() + row_dim, in Compute()
643 b_t_transposed_shape.AddDim(b_t.dim_size(row_dim + 1)); in Compute()
644 b_t_transposed_shape.AddDim(b_t.dim_size(row_dim)); in Compute()
659 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.cc191 int64 row_dim = dim_nums.lhs_batch_dimensions_size(); in RunGemm() local
198 CHECK_NE(row_dim, batch_dim); in RunGemm()
205 CHECK_LT(shape->layout().minor_to_major(row_dim), 2); in RunGemm()
209 int64 output_num_rows = output_shape.dimensions(row_dim); in RunGemm()
234 bool is_row_major = LayoutUtil::Minor(shape.layout(), row_dim) != 0; in RunGemm()
235 bool layout_mismatch = LayoutUtil::Minor(shape.layout(), row_dim) != in RunGemm()
236 LayoutUtil::Minor(output_shape.layout(), row_dim); in RunGemm()
239 shape.dimensions(row_dim + static_cast<int64>(is_row_major)), in RunGemm()
240 shape.dimensions(row_dim + static_cast<int64>(!is_row_major))}; in RunGemm()
244 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.cc257 for (int64_t row_dim = strides.size() - 2; row_dim >= 0; in GetNextInternal() local
258 --row_dim) { in GetNextInternal()
259 strides[row_dim] = in GetNextInternal()
260 strides[row_dim + 1] * t.shape().dim_size(row_dim + 1); in GetNextInternal()