| /external/tensorflow/tensorflow/python/ops/ragged/ |
| D | ragged_gather_ops.py | 40 batch_dims=0, argument
114 def _gather(params, indices, axis, batch_dims): argument
176 def _batch_gather(params, indices, axis, batch_dims): argument
330 def gather_nd(params, indices, batch_dims=0, name=None): argument
|
| D | ragged_gather_op_test.py | 154 batch_dims=0, argument
358 batch_dims=0): argument
|
| D | ragged_dispatch.py | 418 axis=0, batch_dims=0): argument
428 def _ragged_gather_nd_v1(params, indices, name=None, batch_dims=0): argument
|
| /external/tensorflow/tensorflow/python/kernel_tests/ |
| D | banded_triangular_solve_op_test.py | 31 def _verifySolveAllWays(self, x, y, dtypes, batch_dims=None): argument
44 def _verifySolveAllWaysReal(self, x, y, batch_dims=None): argument
47 def _verifySolveAllWaysComplex(self, x, y, batch_dims=None): argument
55 batch_dims=None, argument
|
| D | matrix_triangular_solve_op_test.py | 31 def _verifySolveAllWays(self, x, y, dtypes, batch_dims=None): argument
44 def _verifySolveAllWaysReal(self, x, y, batch_dims=None): argument
47 def _verifySolveAllWaysComplex(self, x, y, batch_dims=None): argument
55 batch_dims=None, argument
|
| D | matrix_solve_op_test.py | 41 def _verifySolve(self, x, y, batch_dims=None): argument
|
| D | matrix_exponential_op_test.py | 222 def _TestRandomSmall(dtype, batch_dims, size): argument
|
| /external/tensorflow/tensorflow/python/kernel_tests/array_ops/ |
| D | gather_op_test.py | 447 def testBatchDims(self, params, indices, batch_dims, expected=None, argument
539 batch_dims, axis, output_shape): argument
574 def _batchNumpyGather(self, params, indices, axis, batch_dims): argument
|
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | gather_scatter_handler.cc | 137 const HloGatherInstruction* gather, absl::Span<const int64> batch_dims, in PartitionIndexOnlyPartition()
181 const HloSharding& output_sharding, absl::Span<const int64> batch_dims, in ParititonPassthroughOperand()
216 const HloSharding& output_sharding, absl::Span<const int64> batch_dims, in ParititonTrivialIndexedOperandDimension()
272 std::vector<int64> batch_dims; in ParititonTrivialIndexedOperandDimension() local
316 const HloSharding& output_sharding, absl::Span<const int64> batch_dims, in PartitionIndexParallelDimensions()
464 absl::Span<const int64> batch_dims, in PartitionGather()
676 std::vector<int64> batch_dims; in HandleGather() local
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | qr_expander.cc | 81 Status House(XlaOp x, XlaOp k, absl::Span<const int64> batch_dims, in House()
190 std::vector<int64> batch_dims(num_batch_dims); in QrBlock() local
295 PrimitiveType type, absl::Span<const int64> batch_dims, XlaOp vs, in CompactWYRepresentation()
370 std::vector<int64> batch_dims(num_batch_dims); in BuildQrDecomposition() local
|
| D | dot_as_convolution_util.h | 43 std::vector<DimNums> batch_dims; member
|
| /external/tensorflow/tensorflow/compiler/xla/client/lib/ |
| D | lu_decomposition.cc | 35 const std::vector<int64> batch_dims( in LuDecomposition() local
|
| D | svd.cc | 119 std::vector<int64> batch_dims(num_batch_dims); in HouseRow() local
184 std::vector<int64> batch_dims(num_batch_dims); in HouseCol() local
258 std::vector<int64> batch_dims(num_batch_dims); in HouseHolderBidiagonalization() local
460 std::vector<int64> batch_dims(num_batch_dims); in OneSidedJacobiUpdate() local
825 std::vector<int64> batch_dims(num_batch_dims); in SVD() local
|
| D | self_adjoint_eig.cc | 120 const std::vector<int64> batch_dims(w_shape.dimensions().begin(), in Update() local
430 std::vector<int64> batch_dims(num_batch_dims); in SelfAdjointEig() local
|
| /external/tensorflow/tensorflow/compiler/tests/ |
| D | matrix_solve_op_test.py | 56 def testSolve(self, n, nrhs, batch_dims, rhs_batch_dims, adjoint): argument
|
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | cusolver_rewriter.cc | 60 std::vector<int64> batch_dims(a_shape.dimensions().begin(), in CreateCholesky() local
|
| /external/tensorflow/tensorflow/core/kernels/ |
| D | gather_op.cc | 92 int32 batch_dims = batch_dims_; in Compute() local
|
| D | gather_op_test.cc | 43 void MakeOp(DataType data_type, DataType index_type, int batch_dims = 0) { in MakeOp()
|
| /external/tensorflow/tensorflow/core/ops/ |
| D | sparse_csr_matrix_ops.cc | 279 ShapeHandle batch_dims; in __anon59a9d6c30702() local
432 ShapeHandle batch_dims; in __anon59a9d6c30a02() local
|
| D | resource_variable_ops.cc | 253 int32 batch_dims; in __anon50d007a70402() local
|
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | gather_op.cc | 158 int batch_dims, xla::XlaOp* gather_output) { in XlaGatherWithBatchDimsOpImpl()
|
| D | softmax_op.cc | 63 std::vector<int64> batch_dims(logits_shape.dims() - 1); in Compile() local
|
| /external/tensorflow/tensorflow/python/ops/ |
| D | array_ops.py | 4831 batch_dims=0): # pylint: disable=g-doc-args argument
5041 batch_dims=0, argument
5071 def _batch_gather(params, indices, batch_dims, axis=None): argument
5196 def gather_nd(params, indices, name=None, batch_dims=0): argument
5359 def gather_nd_v2(params, indices, batch_dims=0, name=None): argument
5366 def batch_gather_nd(params, indices, batch_dims, name=None): argument
|
| D | array_grad.py | 585 def _GetBatchIndices(params_shape, indices, batch_dims): argument
606 def _BatchGatherGrad(params_shape, values, indices, batch_dims, argument
|
| /external/tensorflow/tensorflow/core/grappler/costs/ |
| D | op_level_cost_estimator.h | 211 std::vector<int> batch_dims; member
|