/external/tensorflow/tensorflow/compiler/xla/service/ |
cholesky_expander.cc | 57 TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a)); 58 const int n_dims = a_shape.rank(); 59 const int64 n = ShapeUtil::GetDimension(a_shape, -1); 60 auto major_dims = AsInt64Slice(a_shape.dimensions()) 77 ShapeUtil::MakeShape(a_shape.element_type(), row_shape_dims)); 83 ShapeUtil::MakeShape(a_shape.element_type(), col_shape_dims)); 142 TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a)); 143 const int ndims = a_shape.rank(); 147 a_shape.ToString()); 150 const int64 n = ShapeUtil::GetDimension(a_shape, -1) [all...] |
triangular_solve_expander.cc | 262 TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a)); 263 int64 ndims = a_shape.rank(); 264 int64 n = ShapeUtil::GetDimension(a_shape, -1); 356 TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a)); 358 if (a_shape.rank() != b_shape.rank()) { 362 ShapeUtil::HumanString(a_shape), ShapeUtil::HumanString(b_shape)); 364 const int64 ndims = a_shape.rank(); 373 int64 a_size = a_shape.dimensions(i); 379 ShapeUtil::HumanString(a_shape), ShapeUtil::HumanString(b_shape)); 384 if (ShapeUtil::GetDimension(a_shape, -1) ! [all...] |
hlo_cost_analysis.cc | 555 const Shape& a_shape = hlo->operand(0)->shape(); local 558 int64 elems = a_shape.dimensions(a_shape.dimensions_size() - 1); 568 const Shape& a_shape = hlo->operand(0)->shape(); local 570 int64 elems = a_shape.dimensions(a_shape.dimensions_size() - 1); 571 elems *= ShapeUtil::ElementsIn(a_shape);
|
/external/tensorflow/tensorflow/contrib/factorization/python/kernel_tests/ |
masked_matmul_benchmark.py | 62 def _run_graph(self, a_shape, b_shape, nnz, num_iters, sort=False, 67 a_shape: int list, the shape of the a matrix. 82 mask_shape = [a_shape[0], b_shape[1]] 83 a_shape = a_shape if not transpose_a else [a_shape[1], a_shape[0]] 85 a_var = variables.Variable(random_ops.random_normal(a_shape)) 88 a_ph = array_ops.placeholder(dtypes.float32, shape=a_shape) 110 "cpu nnz:{nnz} a_shape:{a_shape} b_shape:{b_shape} tr_a:{tr_a} [all...] |
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
matmul_op.cc | 53 const TensorShape a_shape = ctx->InputShape(0); variable 58 ctx, TensorShapeUtils::IsMatrix(a_shape), 60 a_shape.DebugString())); 69 a_shape.dim_size(first_index) == b_shape.dim_size(second_index), 71 a_shape.DebugString(), ", In[1]: ",
|
einsum_op.cc | 41 const TensorShape a_shape = ctx->InputShape(0); variable
|
/external/tensorflow/tensorflow/core/kernels/ |
betainc_op.cc | 47 const TensorShape& a_shape = a.shape(); variable 50 if (a_shape.dims() > 0 && b_shape.dims() > 0) { 51 OP_REQUIRES(ctx, a_shape == b_shape, 54 a_shape.DebugString(), " vs. ", b_shape.DebugString())); 56 if (a_shape.dims() > 0 && x_shape.dims() > 0) { 57 OP_REQUIRES(ctx, a_shape == x_shape, 60 a_shape.DebugString(), " vs. ", x_shape.DebugString())); 69 TensorShape merged_shape(a_shape); 76 if (a_shape == b_shape && a_shape == x_shape) [all...] |
sparse_add_op.cc | 32 const Tensor *a_indices, *b_indices, *a_values_t, *b_values_t, *a_shape, variable 65 OP_REQUIRES_OK(ctx, ctx->input("a_shape", &a_shape)); 68 TensorShapeUtils::IsVector(a_shape->shape()) && 72 a_shape->shape().DebugString(), " and ", 75 ctx, a_shape->IsSameSize(*b_shape), 78 a_shape->SummarizeValue(10), " and ", b_shape->SummarizeValue(10))); 79 const auto a_shape_flat = a_shape->flat<int64>(); 81 for (int i = 0; i < a_shape->NumElements(); ++i) { 102 const int num_dims = a_shape->dim_size(0) [all...] |
sparse_tensor_dense_add_op.cc | 36 const Tensor *a_shape, const Tensor *b) { 43 !TensorShapeUtils::IsVector(a_shape->shape())) { 45 "Inputs a_values and a_shape should be vectors " 48 a_shape->shape().DebugString()); 50 if (a_shape->NumElements() != b->dims()) { 52 "Two operands have different ranks; received: ", a_shape->NumElements(), 55 const auto a_shape_flat = a_shape->flat<Index>(); 78 OP_REQUIRES_OK(ctx, ctx->input("a_shape", &a_shape_t));
|
sparse_tensor_dense_matmul_op_test.cc | 27 Node* a_shape, Node* b, bool adjoint_a, 33 .Input(a_shape) 47 Tensor a_shape(DT_INT64, TensorShape({2})); 48 auto a_shape_t = a_shape.vec<int64>(); 66 test::graph::Constant(g, a_values), test::graph::HostConstant(g, a_shape),
|
sparse_tensor_dense_matmul_op.cc | 44 const Tensor* a_shape; variable 48 OP_REQUIRES_OK(ctx, ctx->input("a_shape", &a_shape)); 55 OP_REQUIRES(ctx, TensorShapeUtils::IsVector(a_shape->shape()), 56 errors::InvalidArgument("Tensor 'a_shape' is not a vector")); 59 ctx, a_shape->NumElements() == 2, 60 errors::InvalidArgument("Tensor 'a_shape' must have 2 elements")); 74 ctx, a_indices->shape().dim_size(1) == a_shape->NumElements(), 76 "number of entries in a_shape")); 78 auto a_shape_t = a_shape->vec<int64>() [all...] |
sparse_sparse_binary_op_shared.cc | 128 OP_REQUIRES_OK(ctx, ctx->input("a_shape", &a_shape_t)); 174 const auto a_shape = a_shape_t->flat<int64>(); variable 177 OP_REQUIRES(ctx, a_shape(i) == b_shape(i), 179 a_shape(i), " and ", b_shape(i),
|
/external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
cusolver_rewriter.cc | 57 Shape a_shape = operand->shape(); local 58 int ndim = a_shape.dimensions_size(); 60 int64 n = a_shape.dimensions(ndim - 1); 62 int64 batch_size = std::accumulate(a_shape.dimensions().begin(), 63 a_shape.dimensions().end() - 2, int64{1}, 69 int64 workspace_size; // Number of elements of size a_shape.element_type() 70 switch (a_shape.element_type()) { 101 a_shape.ToString()); 109 SetFortranLayout(&a_shape); 118 {a_shape, [all...] |
/external/tensorflow/tensorflow/compiler/xla/client/lib/ |
qr.cc | 156 TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a)); 157 const int num_dims = a_shape.rank(); 160 a_shape.ToString()); 162 PrimitiveType type = a_shape.element_type(); 164 const int64 m = ShapeUtil::GetDimension(a_shape, -2); 165 const int64 n = ShapeUtil::GetDimension(a_shape, -1); 170 batch_dims[i] = ShapeUtil::GetDimension(a_shape, i); 327 TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a)); 328 const int num_dims = a_shape.rank(); 331 a_shape.ToString()) [all...] |
svd.cc | 117 TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a)); 118 const int64 num_dims = a_shape.rank(); 119 const int64 n = ShapeUtil::GetDimension(a_shape, -1); 126 batch_dims[k] = ShapeUtil::GetDimension(a_shape, k); 183 TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a)); 184 const int64 num_dims = a_shape.rank(); 185 const int64 m = ShapeUtil::GetDimension(a_shape, -2); 192 batch_dims[k] = ShapeUtil::GetDimension(a_shape, k); 261 TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a)); 262 const int64 num_dims = a_shape.rank() 835 Shape a_shape = shape_with_status.ValueOrDie(); local [all...] |
self_adjoint_eig.cc | 80 TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a)); 406 Shape a_shape = shape_with_status.ValueOrDie(); local 407 const int64 num_dims = a_shape.rank(); 411 a_shape.ToString())); 413 PrimitiveType type = a_shape.element_type(); 416 "Type of the input matrix must be float: got %s.", a_shape.ToString())); 419 const int64 m = ShapeUtil::GetDimension(a_shape, -2); 420 const int64 n = ShapeUtil::GetDimension(a_shape, -1); 432 batch_dims[i] = ShapeUtil::GetDimension(a_shape, i);
|
/external/tensorflow/tensorflow/core/ops/ |
sparse_ops.cc | 32 TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 1, &unused)); // a_shape 65 .Input("a_shape: int64") 76 ShapeHandle a_shape; 77 TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 1, &a_shape)); 79 0, c->Matrix(InferenceContext::kUnknownDim, c->Dim(a_shape, 0))); 81 c->set_output(2, a_shape); 88 .Input("a_shape: int64") 99 ShapeHandle a_shape; 102 TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensor(2, &a_shape)); 103 TF_RETURN_IF_ERROR(c->WithRank(a_shape, 2, &a_shape)) [all...] |
math_ops.cc | 128 ShapeHandle a_shape; 130 TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 2, &a_shape)); 138 DimensionHandle output_rows = c->Dim(a_shape, adj_x ? -1 : -2); 145 TF_RETURN_IF_ERROR(c->Subshape(a_shape, 0, -2, &a_batch_dims)); 151 TF_RETURN_IF_ERROR(c->Merge(c->Dim(a_shape, adj_x ? -2 : -1), [all...] |
/external/tensorflow/tensorflow/python/ops/ |
linalg_grad.py | 355 a_shape = a.get_shape().with_rank_at_least(2) 361 grad_a.set_shape(a_shape) 373 m = a_shape.dims[-2].merge_with(grad_u_shape[-2]) 374 n = a_shape.dims[-1].merge_with(grad_v_shape[-2]) 375 batch_shape = a_shape[:-2].merge_with(grad_u_shape[:-2]).merge_with( 377 a_shape = batch_shape.concatenate([m, n]) 379 m = a_shape.dims[-2].value 380 n = a_shape.dims[-1].value 466 grad_a.set_shape(a_shape)
|
sparse_grad.py | 80 (a_indices, a_values, a_shape, b_indices, b_values, b_shape, thresh) 95 # (a_indices, a_values, a_shape, b_indices, b_values, b_shape, thresh) 166 a_indices, a_values, a_shape = op.inputs[:3] 182 a_indices, a_values, a_shape, grad, adjoint_a=not adj_a) 200 # gradients w.r.t. (a_indices, a_values, a_shape, b)
|
math_ops.py | [all...] |
/external/tensorflow/tensorflow/compiler/tests/ |
matrix_triangular_solve_op_test.py | 101 for dtype, (a_shape, b_shape) in tuples: 102 n = a_shape[-1] 103 a = np.tril(rng.rand(*a_shape) - 0.5) / (2.0 * n) + np.eye(n)
|
/external/tensorflow/tensorflow/python/kernel_tests/ |
batch_matmul_op_test.py | 109 def CompareNonEmpty(self, a_shape, b_shape): 111 self._rand(a_shape, dtype), 126 def CompareEmpty(self, a_shape, b_shape): 128 np.zeros(a_shape).astype(dtype),
|
tensordot_op_test.py | 146 a_shape = np.random.random_integers(1, _MAXDIM, rank_a_) 152 a_shape[a_dims[i]] = shared_shape[i] 156 size=np.prod(a_shape)).reshape(a_shape).astype(dtype_)
|
/external/tensorflow/tensorflow/python/framework/ |
common_shapes.py | 106 a_shape = op.inputs[0].get_shape().with_rank(2) 110 output_rows = a_shape[1] if transpose_a else a_shape[0] 112 inner_a = a_shape[0] if transpose_a else a_shape[1]
|