Searched refs:row_dim (Results 1 – 6 of 6) sorted by relevance
/external/tensorflow/tensorflow/core/kernels/sparse/ |
D | sparse_mat_mul_op.cc | 138 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 …]
|
D | add_op.cc | 280 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()
|
D | mat_mul_op.cc | 515 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()
|
D | sparse_cholesky_op.cc | 244 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/ |
D | gemm_thunk.cc | 181 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/ |
D | dense_to_sparse_batch_dataset_op.cc | 251 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()
|