Searched refs:reduction_dim (Results 1 – 8 of 8) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/client/lib/ |
| D | approx_topk.cc | 112 int64_t reduction_dim, in AggregateToTopKBuilder() argument
128 reduction_computation, {reduction_dim}); in AggregateToTopKBuilder()
130 op_shape.mutable_dimensions()[reduction_dim] = 1; in AggregateToTopKBuilder()
138 auto sorted_results = Sort(operands, comparator, reduction_dim); in AggregateToTopKBuilder()
145 slice_limit_indices[reduction_dim] = top_k; in AggregateToTopKBuilder()
159 int64_t reduction_dim, const XlaComputation& comparator, in ApproxTopK() argument
177 if (reduction_dim < 0 || reduction_dim >= rank) { in ApproxTopK()
186 uint64_t n = operands_shapes[0].dimensions(reduction_dim); in ApproxTopK()
191 reduction_dim, comparator); in ApproxTopK()
210 reduction_dim, comparator); in ApproxTopK()
[all …]
|
| D | approx_topk.h | 55 int64_t reduction_dim, const XlaComputation& comparator,
62 int64_t reduction_dim,
|
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | approx_topk_op.cc | 70 int64_t reduction_dim = reduction_dim_; in Compile() local
71 if (reduction_dim < 0) { in Compile()
73 reduction_dim += op_shape.dimensions_size(); in Compile()
87 reduction_dim); in Compile()
90 reduction_dim, comparator, recall_target_, aggregate_to_topk_, in Compile()
100 int64_t reduction_dim, const xla::XlaComputation& comparator,
123 int64_t reduction_dim, const xla::XlaComputation& comparator, in ApproxTopKFn() argument
126 return xla::ApproxTopK(builder, operands, init_values, top_k, reduction_dim, in ApproxTopKFn()
141 int64_t reduction_dim, const xla::XlaComputation& comparator, in ApproxTopKFn() argument
145 builder, operands, init_values, top_k, reduction_dim, comparator, in ApproxTopKFn()
|
| /external/eigen/unsupported/Eigen/CXX11/src/Tensor/ |
| D | TensorReductionSycl.h | 250 enum class reduction_dim { inner_most, outer_most };
252 template <typename Evaluator, typename OpType, typename PannelParameters, reduction_dim rt>
287 …Index global_offset = rt == reduction_dim::outer_most ? globalPId + (globalRId * num_coeffs_to_pre…
293 …rt == reduction_dim::outer_most ? num_coeffs_to_preserve * per_thread_local_stride : per_thread_lo…
302 …Index pLocalThreadId = rt == reduction_dim::outer_most ? linearLocalThreadId % PannelParameters::L…
304 …Index rLocalThreadId = rt == reduction_dim::outer_most ? linearLocalThreadId / PannelParameters::L…
306 …const Index pGroupId = rt == reduction_dim::outer_most ? itemID.get_group(0) % preserve_elements_n…
308 …const Index rGroupId = rt == reduction_dim::outer_most ? itemID.get_group(0) / preserve_elements_n…
323 if (rt == reduction_dim::inner_most) {
334 if (rt == reduction_dim::inner_most) {
[all …]
|
| /external/tensorflow/tensorflow/python/ops/ |
| D | nn_grad.py | 1153 reduction_dim = op.get_attr("reduction_dimension")
1154 if reduction_dim < 0:
1155 reduction_dim = rank + reduction_dim
1158 if d == reduction_dim:
|
| /external/tensorflow/tensorflow/core/api_def/base_api/ |
| D | api_def_ApproxTopK.pbtxt | 46 `input[reduction_dim]` for evaluating the recall. This option is useful when
|
| /external/tensorflow/tensorflow/compiler/xla/python/xla_extension/ |
| D | ops.pyi | 83 reduction_dim: int,
93 reduction_dim: int,
|
| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | xla_client_test.py | 1637 reduction_dim=1,
|