| /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]
|
