Home
last modified time | relevance | path

Searched defs:arg_shapes (Results 1 – 21 of 21) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/
H A Dcompile_mlir_util.cc116 mlir::ModuleOp module, llvm::ArrayRef<TensorOrResourceShape> arg_shapes, in GetXlaInputShapes()
326 Status RefineShapes(llvm::ArrayRef<TensorOrResourceShape> arg_shapes, in RefineShapes()
554 llvm::ArrayRef<TensorOrResourceShape> arg_shapes) { in CompileMlirSetup()
567 llvm::ArrayRef<TensorOrResourceShape> arg_shapes, in BuildHloFromTf()
615 mlir::ModuleOp module_op, llvm::ArrayRef<TensorOrResourceShape> arg_shapes, in PopulateResultIOInfo()
636 mlir::ModuleOp module_op, llvm::ArrayRef<TensorOrResourceShape> arg_shapes, in CompileMlirToXlaHlo()
666 llvm::StringRef mlir_module_string, llvm::ArrayRef<TensorShape> arg_shapes, in CompileSerializedMlirToXlaHlo()
767 llvm::SmallVector<TensorOrResourceShape, 4>& arg_shapes) { in CompileGraphSetup()
809 llvm::SmallVector<TensorOrResourceShape, 4> arg_shapes; in BuildHloFromModule() local
828 llvm::SmallVector<TensorOrResourceShape, 4> arg_shapes; in CompileGraphToXlaHlo() local
H A Dtf_xla_mlir_translate.cc133 llvm::SmallVectorImpl<TensorOrResourceShape>& arg_shapes) { in ParseArgumentShapes()
255 mlir::ModuleOp module_op, llvm::ArrayRef<TensorOrResourceShape> arg_shapes, in CompileMlirToXlaHloViaBuilder()
313 llvm::SmallVector<TensorOrResourceShape, 4> arg_shapes; in MlirTfToHloTextTranslateFunctionImpl() local
/aosp_15_r20/external/tensorflow/tensorflow/core/tpu/
H A Dtpu_compile.cc165 const std::vector<PartialTensorShape>& arg_shapes, in OptimizeGraph()
253 const std::vector<TensorShape>& arg_shapes, in BuildComputationArgumentDescriptions()
340 const std::vector<PartialTensorShape>& arg_shapes, Graph* graph, in RunShapeInferenceOnComputation()
374 const std::vector<TensorShape>& arg_shapes, in CompileTFFunctionToHlo()
453 absl::Span<const TensorShape> arg_shapes, in GetShardingInfo()
/aosp_15_r20/external/tensorflow/tensorflow/core/tpu/kernels/
H A Dtpu_compile_op_support.cc168 std::vector<Shape> arg_shapes; in GetPerDeviceShape() local
345 const std::vector<TensorShape>& arg_shapes) { in CreateTpuCompilationRequest()
449 std::vector<TensorShape>* arg_shapes) { in ComputeArgumentShapes()
H A Dtpu_compile_op_common.cc185 std::vector<TensorShape> arg_shapes; in CompileLocallyAndFillHostCacheInternal() local
468 const std::vector<TensorShape>& arg_shapes, in RegisterXLAFingerprints()
H A Dtpu_compile.proto45 repeated TensorShapeProto arg_shapes = 6; field
H A Dtpu_compile_op_impl.cc31 const std::vector<TensorShape>& arg_shapes, in Compile()
H A Dtpu_functional_ops.cc402 std::map<int, InferredShape>& arg_shapes, EdgeShapes& tpu_input_shapes, in GetInputOutputInfo()
531 std::vector<TensorShape> arg_shapes; in MaybeRegisterFingerprint() local
1997 std::map<int, InferredShape>& arg_shapes, in InferShapesWithResourceVar()
2197 std::map<int, InferredShape> arg_shapes; in OptimizeTpuInputOutputTensors() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/jit/
H A Dshape_inference.cc46 const std::map<int, InferredShape>& arg_shapes, in PropagateShapes()
243 Status InferShapes(Graph* graph, const std::map<int, InferredShape>& arg_shapes, in InferShapes()
H A Dencapsulate_util.cc323 std::map<int, InferredShape> arg_shapes; in PerformStaticShapeInferenceBeforeEncapsulation() local
/aosp_15_r20/external/tensorflow/tensorflow/core/kernels/data/
H A Dmap_defun_op.cc48 const std::vector<TensorShape> arg_shapes; member
279 std::vector<TensorShape> arg_shapes; in SetupArgs() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dshape_inference.cc364 absl::Span<const Shape* const> arg_shapes, const int64_t dimension) { in InferConcatOpShape()
1115 absl::Span<const Shape* const> arg_shapes, const ProgramShape& to_apply, in InferMapShape()
2217 absl::Span<const Shape* const> arg_shapes, in InferReduceShape()
3298 absl::Span<const Shape* const> arg_shapes, const ProgramShape& to_apply) { in InferCallShape()
3658 absl::Span<const Shape* const> arg_shapes, in InferScatterShape()
H A Dhlo_parser.cc1782 absl::InlinedVector<const Shape*, 2> arg_shapes; in CreateInstruction() local
1801 absl::InlinedVector<const Shape*, 2> arg_shapes; in CreateInstruction() local
1935 absl::InlinedVector<const Shape*, 2> arg_shapes; in CreateInstruction() local
2143 absl::InlinedVector<const Shape*, 2> arg_shapes; in CreateInstruction() local
2168 absl::InlinedVector<const Shape*, 2> arg_shapes; in CreateInstruction() local
2199 absl::InlinedVector<const Shape*, 2> arg_shapes; in CreateInstruction() local
2890 absl::InlinedVector<const Shape*, 3> arg_shapes; in CreateInstruction() local
H A Dhlo_verifier.cc1663 absl::InlinedVector<const Shape*, 3> arg_shapes; in HandleScatter() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2xla/
H A Dxla_jit_compiled_cpu_function.cc111 std::vector<const xla::Shape*> arg_shapes; in Compile() local
H A Dxla_compiler.cc1050 std::vector<xla::Shape> arg_shapes(input_to_args->size()); in BuildArguments() local
/aosp_15_r20/external/tensorflow/tensorflow/core/tpu/graph_rewrite/
H A Ddistributed_tpu_rewrite_pass.cc1907 const ParameterInfo& params_info, std::vector<InferredShape>* arg_shapes, in GetArgAndRetvalShapes()
2127 const std::vector<InferredShape>& arg_shapes, in AssignArgsAndRetvalsToCores()
2464 const Node& replicate_node, const std::vector<InferredShape>& arg_shapes, in BuildDynamicShapeNodes()
2595 const std::vector<InferredShape>& arg_shapes, in BuildCompileNode()
3158 const std::vector<InferredShape>& arg_shapes, in CreateOrGetPerHostVariableCopy()
3243 const std::vector<InferredShape>& arg_shapes, in BuildExecuteNodes()
4754 std::vector<InferredShape> arg_shapes, retval_shapes; in RewriteTPUReplicateNode() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/pjrt/
H A Dutils.cc78 std::vector<Shape> arg_shapes; in GetShardedProgramShapes() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/xla/
H A Dmlir_hlo_to_hlo.cc2322 llvm::SmallVectorImpl<xla::Shape>* arg_shapes, in SetEntryTupleShapesAndLeafReplication()
2359 llvm::SmallVectorImpl<xla::Shape>* arg_shapes) { in SetEntryTupleShardings()
2399 llvm::SmallVector<xla::Shape, 4> arg_shapes; in LowerBasicBlockAsFunction() local
2426 llvm::SmallVector<xla::Shape, 4> arg_shapes; in LowerBasicBlockAsFunction() local
/aosp_15_r20/external/tensorflow/tensorflow/python/framework/
H A Dfunction.py946 arg_shapes=None, argument
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
H A Dshape_inference.cc2684 ArrayRef<ArrayRef<int64_t>> arg_shapes, in InferShapeForFunction()