Home
last modified time | relevance | path

Searched defs:output_rank (Results 1 – 25 of 37) sorted by relevance

12

/external/tensorflow/tensorflow/dtensor/mlir/expansions/
Dslice_spmd_expander.cc387 const int output_rank, in ApplyNewAndShrinkMasksToLayout()
616 const int output_rank = ValueRank(strided_slice_op.output()); in ExpandOp() local
656 const int output_rank = ValueRank(strided_slice_op.output()); in ComputeLayoutForward() local
682 const int output_rank = ValueRank(strided_slice_op.output()); in ComputeLayoutBackward() local
873 const int output_rank = global_output_shape.size(); in ExpandOp() local
940 const int output_rank = global_output_shape.size(); in ComputeLayoutForward() local
967 const int output_rank = global_output_shape.size(); in ComputeLayoutBackward() local
Dsegmentation_spmd_expander.cc38 const int output_rank = ValueRank(unsorted_segmented_sum.output()); in ComputeLayoutForward() local
Delementwise_spmd_expander.cc129 const int output_rank = ValueRank(op->getOpResult(0)); in ComputeLayoutForward() local
Dmeta_spmd_expander.cc136 const int output_rank = ValueRank(pack.output()); in ExpandOp() local
1235 const int output_rank = ValueRank(op->getResult(0)); in ComputeLayoutForward() local
/external/tensorflow/tensorflow/core/kernels/
Dreshape_util_gpu.cu.cc31 const Tindex nnz, const Tindex input_rank, const Tindex output_rank, in ReshapeSparseTensorKernel()
64 const int64 output_rank = output_shape.dims(); in operator ()() local
Dreshape_util.cc49 const int64_t output_rank = output_shape.dims(); in operator ()() local
103 const int64_t output_rank = target_shape_in.NumElements(); in ReshapeSparseTensor() local
Ddeserialize_sparse_variant_op.cc155 const int output_rank = output_shape->NumElements(); in Compute() local
/external/executorch/kernels/portable/cpu/
Dop_expand_copy.cpp74 size_t output_rank = 0; in expand_copy_out() local
/external/tensorflow/tensorflow/compiler/xla/service/
Dgather_expander.cc256 int64_t output_rank) { in PermuteBatchAndOffsetDims()
357 int64_t output_rank = output_shape.dimensions_size(); in ExpandInstruction() local
Dgather_simplifier.cc113 auto output_rank = static_cast<int64_t>(start_indices_dims.size() + in ExpandInstruction() local
Dinstruction_fusion.cc212 int64_t output_rank = 0; in EffectivelyAtMostUnary() local
Dhlo_evaluator.cc2153 const int64_t output_rank = output_shape.rank(); in CheckParameters() local
2198 int64_t output_rank = output_shape.dimensions_size(); in IterationSpaceForOutputBatchIndices() local
2215 int64_t output_rank, absl::Span<const int64_t> slice_sizes, in IterationSpaceForOutputOffsetIndices()
/external/tensorflow/tensorflow/core/ops/
Dset_ops.cc53 DimensionHandle output_rank; in __anonbfa34e750102() local
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dbincount_op.cc48 auto output_rank = output_shape_param.rank(); in Compile() local
/external/tensorflow/tensorflow/lite/kernels/
Dgather_nd.cc93 const int output_rank = indices_rank + params_rank - indices_nd - 1; in Prepare() local
Dembedding_lookup_sparse.cc170 const int output_rank = (lookup_rank - 1) + (embedding_rank - 1); in Eval() local
Dunsorted_segment.cc78 const int output_rank = data_rank - segment_ids_rank + 1; in ResizeOutputTensor() local
Dbatch_matmul.cc102 bool adj_x, bool adj_y, int output_rank, in ResizeOutputTensor()
362 const int output_rank = std::max(lhs_rank, rhs_rank); in Prepare() local
/external/tflite-support/tensorflow_lite_support/custom_ops/kernel/
Dunsorted_segment.cc110 const int output_rank = data_rank - segment_ids_rank + 1; in ResizeOutputTensor() local
/external/tensorflow/tensorflow/compiler/xla/service/llvm_ir/
Dir_array.cc384 int64_t output_rank = shape.rank(); in SourceIndexOfBroadcast() local
/external/tensorflow/tensorflow/compiler/tf2tensorrt/convert/ops/
Deinsum.cc527 int output_rank = desc_a.b + desc_a.f + desc_b.f; in ShuffleEinsumOutput() local
/external/executorch/kernels/portable/cpu/util/
Dcopy_ops_util.cpp187 size_t* output_rank) { in get_expand_copy_out_target_size()
/external/tensorflow/tensorflow/compiler/xla/client/lib/
Dmatrix.cc413 const int64_t output_rank = output_config.size(); in Einsum() local
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
Dlegalize_hlo.cc3010 int output_rank = result_type.getRank(); in CanonicalizeOffset() local
3274 int64_t output_rank = output.getType().cast<ShapedType>().getRank(); in IsTFStyleBroadcast() local
/external/tensorflow/tensorflow/compiler/xla/
Dshape_util.cc1562 int64_t output_rank = output_shape.rank(); in AlignLayouts() local

12