Home
last modified time | relevance | path

Searched defs:global_shape (Results 1 – 13 of 13) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/dtensor/cc/
H A Dsave_restore_util.cc42 absl::Span<const int64_t> global_shape, Layout layout) { in BuildSliceSpecDeviceMap()
166 absl::Span<const int64_t> global_shape) { in SliceSpecOnDevice()
H A Ddtensor_sharded_prefix_kernel.cc102 std::vector<int64_t> global_shape = in Compute() local
H A Ddtensor_device_util.h317 const std::vector<int64_t> global_shape() const { in global_shape() function
H A Dtensor_layout.cc975 std::vector<int64_t> global_shape; in GlobalShapeFromLocalShape() local
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/
H A Dshape_utils.cc45 auto global_shape = layout_op.global_shape(); in ExtractGlobalInputShape() local
81 auto global_shape = layout_op.global_shape(); in ExtractGlobalOutputShape() local
H A Dlayout_propagation_v2.cc1181 mlir::TF::ShapeAttr global_shape = mlir::TF::ShapeAttr::get( in InsertRelayoutForWhileLoops() local
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/expansions/
H A Dcontrol_flow_spmd_expander.cc75 const auto global_shape = result_layout_op.global_shape(); in ExpandOp() local
H A Dconv_spmd_expander.cc291 llvm::SmallVector<int64_t, 4> global_shape; in HandleConvBackpropInput() local
H A Dsoftmax_spmd_expander.cc187 StatusOr<Layout> GetBroadcastedLayout(llvm::ArrayRef<int64_t> global_shape, in GetBroadcastedLayout()
H A Dsave_restore_spmd_expander.cc578 std::vector<int64_t> global_shape = std::get<0>(it.value()); in ExpandRestoreV2OpHelper() local
H A Dslice_spmd_expander.cc55 llvm::ArrayRef<int64_t>* global_shape = nullptr) { in VerifySliceLayout()
H A Dmeta_spmd_expander.cc1198 auto global_shape = builder.create<mlir::TF::MulOp>( in ExpandOp() local
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/python/
H A Dd_variable.py41 global_shape, argument