Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/dtensor/mlir/expansions/
Dexpanddims_spmd_expander.cc47 ExtractConstIntFromValue(expand_dims_op.dim())); in ExpandOp()
81 ExtractConstIntFromValue(expand_dims_op.dim())); in ComputeLayoutForward()
113 ExtractConstIntFromValue(expand_dims_op.dim())); in ComputeLayoutBackward()
Dargmax_spmd_expander.cc49 ExtractConstIntFromValue(argmax_op.dimension())); in ComputeResultLayout()
68 ExtractConstIntFromValue(argmax_op.dimension())); in ExpandOp()
145 ExtractConstIntFromValue(argmax_op.dimension())); in ComputeLayoutBackward()
Dconcat_spmd_expander.cc37 ExtractConstIntFromValue(concat_dim_operand)); in VerifyConcatLayout()
54 ExtractConstIntFromValue(concat_dim_operand)); in ReduceForConcatOutputLayout()
Dgather_spmd_expander.cc37 TF_ASSIGN_OR_RETURN(int64_t axis, ExtractConstIntFromValue(gather_op.axis())); in ExpandOp()
173 TF_ASSIGN_OR_RETURN(int64_t axis, ExtractConstIntFromValue(gather_op.axis())); in ComputeLayoutForward()
252 ExtractConstIntFromValue(gather_op.axis())); in ComputeLayoutBackward()
Dcumsum_spmd_expander.cc46 ExtractConstIntFromValue(cumsum.axis())); in GetAxisDimension()
Dsplit_spmd_expander.cc71 ExtractConstIntFromValue(split_dim_value)); in GetAdjustedSplitDim()
Dmeta_spmd_expander.cc1008 const auto depth = ExtractConstIntFromValue(one_hot_op.depth()); in ExpandOp()
/external/tensorflow/tensorflow/dtensor/mlir/
Dvalue_utils.h53 StatusOr<int64_t> ExtractConstIntFromValue(mlir::Value value);
Dvalue_utils.cc122 StatusOr<int64_t> ExtractConstIntFromValue(mlir::Value value) { in ExtractConstIntFromValue() function