Home
last modified time | relevance | path

Searched defs:gather_indices (Results 1 – 7 of 7) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
H A Dresampler_ops.cc141 XlaOp Gather2by2Neighbors(xla::XlaBuilder* b, XlaOp data, XlaOp gather_indices, in Gather2by2Neighbors()
252 XlaOp gather_indices, XlaOp warp, in CalculateGradData()
363 XlaOp gather_indices, XlaOp data, in CalculateGradWarp()
521 auto gather_indices = ConcatenateIota(ctx->builder(), top_left, warp_shape); in Compile() local
654 auto gather_indices = ConcatenateIota(ctx->builder(), top_left, warp_shape); in Compile() local
/aosp_15_r20/external/pytorch/torch/csrc/jit/passes/onnx/
H A Dpeephole.cpp336 Node* gather_indices = b->owningGraph()->create(onnx::Constant, 1); in pushPackingPastRnn() local
480 Node* gather_indices = graph->create(onnx::Constant, 1); in fixDefaultRNNState() local
H A Dpreprocess_for_onnx.cpp204 Node* gather_indices = b->owningGraph()->create(onnx::Constant, 1); in fuseListAndListUnpack() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/tests/
H A Dgather_test.py157 def _benchmarkGather(self, name, axis, gather_indices, use_xla_jit): argument
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/lib/
H A Dslicing.cc215 XlaOp gather_indices = ConcatInDim(builder, to_concat, input_shape.rank()); in TorchGather() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2tensorrt/convert/ops/
H A Dlayer_utils.h449 nvinfer1::ITensor* gather_indices = (*indices_result)->getOutput(0); variable
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Ddynamic_padder.cc710 HloInstruction* gather_indices = reshape->AddInstruction( in RewriteDynamicReshapeCombineInput() local