Home
last modified time | relevance | path

Searched defs:index_shape (Results 1 – 17 of 17) sorted by relevance

/external/tensorflow/tensorflow/core/tpu/kernels/xla/
Dget_item_op.cc37 const TensorShape& index_shape = ctx->InputShape(1); in Compile() local
/external/tensorflow/tensorflow/lite/kernels/
Dembedding_lookup_test.cc41 BaseEmbeddingLookupOpModel(std::initializer_list<int> index_shape, in BaseEmbeddingLookupOpModel()
90 HybridEmbeddingLookupOpModel(std::initializer_list<int> index_shape, in HybridEmbeddingLookupOpModel()
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Ddynamic_slice_ops.cc51 const TensorShape index_shape = ctx->InputShape("indices"); in Compile() local
Dsplit_op.cc106 const TensorShape index_shape = ctx->InputShape(2); in Compile() local
/external/tensorflow/tensorflow/compiler/xla/service/
Dscatter_expander.cc136 const Shape& index_shape = index_vector->shape(); in ExpandIndexVectorIntoOperandSpace() local
Dgather_expander.cc114 const Shape& index_shape = index_vector->shape(); in ExpandIndexVectorIntoOperandSpace() local
Dshape_inference.cc2539 for (const Shape& index_shape : start_index_shapes) { in InferDynamicSliceShape() local
2646 for (const Shape& index_shape : start_index_shapes) { in InferDynamicUpdateSliceShape() local
Dhlo_alias_analysis_test.cc1068 Shape index_shape = ShapeUtil::MakeShape(S32, {}); in TEST_F() local
Dalgebraic_simplifier.cc2530 const Shape& index_shape = gather->operand(1)->shape(); in HandleGather() local
/external/tensorflow/tensorflow/core/kernels/batching_util/
Dbatch_resource_base.cc676 TensorShape index_shape({0, 3}); in ProcessBatch() local
708 const TensorShape index_shape({batch.num_tasks(), 3}); in EmitIndexTensor() local
/external/tensorflow/tensorflow/lite/kernels/parse_example/
Dparse_example.cc495 TfLiteIntArray* index_shape = TfLiteIntArrayCreate(2); in FastParseExampleLite() local
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner_util.cc1021 auto index_shape = ShapeUtil::ChangeElementType(valid_slice->shape(), S32); in ExchangeHaloAndGetValidData() local
Dspmd_partitioner.cc489 auto index_shape = ShapeUtil::ChangeElementType(shape, S32); in PadWithValue() local
/external/tensorflow/tensorflow/compiler/tf2xla/
Dxla_compiler_test.cc1618 xla::Shape index_shape; in TEST_F() local
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/ir/
Dtf_ops_a_m.cc1979 ArrayRef<int64_t> index_shape = index_ty.getShape(); in Verify() local
/external/tensorflow/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/IR/
Dhlo_ops.cc3001 auto index_shape = index_type.getShape().vec(); in fold() local
/external/tensorflow/tensorflow/lite/delegates/nnapi/
Dnnapi_delegate_test.cc4752 BaseEmbeddingLookupOpModel(std::initializer_list<int> index_shape, in BaseEmbeddingLookupOpModel()