Home
last modified time | relevance | path

Searched defs:iota_shape (Results 1 – 9 of 9) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/client/lib/
Dsorting.cc48 Shape iota_shape = in TopK() local
90 Shape iota_shape = in TopKWithPartitions() local
Darithmetic.cc157 auto iota_shape = input_shape; in ArgMinMax() local
Dmatrix.cc267 Shape iota_shape = x_shape; in EinsumDiagonalMask() local
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dmatrix_band_part_op.cc64 xla::Shape iota_shape = xla::ShapeUtil::MakeShape(index_xla_type, {m, n}); in Compile() local
Dwhere_op.cc44 auto iota_shape = input_shape.ValueOrDie(); in Compile() local
Dunique_op.cc151 auto iota_shape = input_shape; in DataOutputFastPath() local
/external/tensorflow/tensorflow/compiler/xla/service/
Dstable_sort_expander.cc64 Shape iota_shape = sort->operand(0)->shape(); in ExpandInstruction() local
/external/tensorflow/tensorflow/compiler/tf2xla/
Dxla_helpers.cc94 xla::Shape iota_shape; in OneHot() local
/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/
Dlegalize_tf.cc5934 auto iota_shape = llvm::to_vector<4>(batch_dims); in QRBlock() local
6017 auto iota_shape = llvm::to_vector<4>(batch_dims); in ComputeWYRepresentation() local