Searched refs:iota_op (Results 1 – 3 of 3) sorted by relevance
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| D | legalize_hlo.cc | 713 mhlo::IotaOp iota_op, ArrayRef<Value> args, in matchAndRewrite() argument
716 iota_op.getType().dyn_cast_or_null<RankedTensorType>(); in matchAndRewrite()
719 const uint64_t dimension = iota_op.iota_dimension(); in matchAndRewrite()
736 Value start_op = rewriter.create<TF::ConstOp>(iota_op.getLoc(), start); in matchAndRewrite()
737 Value limit_op = rewriter.create<TF::ConstOp>(iota_op.getLoc(), limit); in matchAndRewrite()
738 Value delta_op = rewriter.create<TF::ConstOp>(iota_op.getLoc(), delta); in matchAndRewrite()
739 Value result = rewriter.create<TF::RangeOp>(iota_op.getLoc(), range_type, in matchAndRewrite()
744 reshape_shape[iota_op.iota_dimension()] = type.getShape()[dimension]; in matchAndRewrite()
747 iota_op.getLoc(), rewriter.getI64TensorAttr(reshape_shape)); in matchAndRewrite()
748 result = rewriter.create<TF::ReshapeOp>(iota_op.getLoc(), reshape_type, in matchAndRewrite()
[all …]
|
| /external/tensorflow/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/ |
| D | legalize_to_linalg.cc | 839 OpTy iota_op, ArrayRef<Value> args, in matchAndRewrite() argument
841 ShapedType result_shaped_type = GetHloOpResultType<isLHLO>(iota_op); in matchAndRewrite()
850 Location loc = iota_op.getLoc(); in matchAndRewrite()
854 GetResultValue<isLHLO>(iota_op)); in matchAndRewrite()
869 nested_loc, ivs[iota_op.iota_dimension()], in matchAndRewrite()
879 rewriter.replaceOp(iota_op, llvm::None); in matchAndRewrite()
881 rewriter.replaceOp(iota_op, linalg_op.result_tensors()); in matchAndRewrite()
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
| D | legalize_tf.cc | 4752 Value iota_op = rewriter.create<mhlo::IotaOp>( in matchAndRewrite() local
4758 op.getLoc(), llvm::ArrayRef<Value>{op.input(), iota_op}, last_dim_index, in matchAndRewrite()
|