Home
last modified time | relevance | path

Searched refs:GetGlobalShapeOfValueFromDTensorLayout (Results 1 – 9 of 9) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/expansions/
H A Dscatter_spmd_expander.cc97 GetGlobalShapeOfValueFromDTensorLayout(scatter_op.tensor())); in TensorScatterOpExpand()
100 GetGlobalShapeOfValueFromDTensorLayout(scatter_op.indices())); in TensorScatterOpExpand()
103 GetGlobalShapeOfValueFromDTensorLayout(scatter_op.updates())); in TensorScatterOpExpand()
H A Dslice_spmd_expander.cc589 GetGlobalShapeOfValueFromDTensorLayout(strided_slice_op.input())); in ExpandOp()
724 GetGlobalShapeOfValueFromDTensorLayout(strided_slice_op.input())); in ExpandOp()
870 GetGlobalShapeOfValueFromDTensorLayout(strided_slice_grad_op.output())); in ExpandOp()
H A Dexpanddims_spmd_expander.cc42 GetGlobalShapeOfValueFromDTensorLayout(expand_dims_op.output())); in ExpandOp()
H A Dbroadcast_to_spmd_expander.cc55 GetGlobalShapeOfValueFromDTensorLayout(broadcast_op.input())); in ExpandOp()
H A Dsoftmax_spmd_expander.cc421 GetGlobalShapeOfValueFromDTensorLayout(op->getOperand(0))); in MaybeRelayoutInputs()
440 GetGlobalShapeOfValueFromDTensorLayout(op->getOperand(1))); in MaybeRelayoutInputs()
H A Dsave_restore_spmd_expander.cc281 GetGlobalShapeOfValueFromDTensorLayout(original_save.tensor_names())); in ConditionalSave()
412 GetGlobalShapeOfValueFromDTensorLayout(it.value())); in ExpandSaveV2Op()
H A Dgather_spmd_expander.cc381 GetGlobalShapeOfValueFromDTensorLayout(gather_op.indices())); in ExpandOp()
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/
H A Dshape_utils.h43 StatusOr<llvm::ArrayRef<int64_t>> GetGlobalShapeOfValueFromDTensorLayout(
H A Dshape_utils.cc246 StatusOr<llvm::ArrayRef<int64_t>> GetGlobalShapeOfValueFromDTensorLayout( in GetGlobalShapeOfValueFromDTensorLayout() function