| 70 | explicit RangeOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} |
| 71 | |
| 72 | void Compile(XlaOpKernelContext* ctx) override { |
| 73 | const TensorShape start_in_shape = ctx->InputShape(0); |
| 74 | const TensorShape limit_in_shape = ctx->InputShape(1); |
| 75 | const TensorShape delta_in_shape = ctx->InputShape(2); |
| 76 | OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(start_in_shape), |
| 77 | errors::InvalidArgument("start must be a scalar, not shape ", |
| 78 | start_in_shape.DebugString())); |
| 79 | OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(limit_in_shape), |
| 80 | errors::InvalidArgument("limit must be a scalar, not shape ", |
| 81 | limit_in_shape.DebugString())); |
| 82 | OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(delta_in_shape), |
| 83 | errors::InvalidArgument("delta must be a scalar, not shape ", |
| 84 | delta_in_shape.DebugString())); |
| 85 | xla::Literal start, limit, delta; |
| 86 | OP_REQUIRES_OK(ctx, ctx->ConstantInput(0, &start)); |
| 87 | OP_REQUIRES_OK(ctx, ctx->ConstantInput(1, &limit)); |
| 88 | OP_REQUIRES_OK(ctx, ctx->ConstantInput(2, &delta)); |
| 89 | |
| 90 | DataType type = input_type(0); |
| 91 | xla::StatusOr<xla::XlaOp> output; |
| 92 | switch (type) { |
| 93 | case DT_INT32: |
| 94 | output = CreateRangeTensor<int32>(start, limit, delta, ctx->builder()); |
| 95 | break; |
| 96 | case DT_INT64: |
| 97 | output = CreateRangeTensor<int64>(start, limit, delta, ctx->builder()); |
| 98 | break; |
| 99 | case DT_FLOAT: |
| 100 | output = CreateRangeTensor<float>(start, limit, delta, ctx->builder()); |
| 101 | break; |
| 102 | case DT_DOUBLE: |
| 103 | output = CreateRangeTensor<double>(start, limit, delta, ctx->builder()); |
| 104 | break; |
| 105 | default: |
| 106 | output = errors::InvalidArgument("Invalid type for Range ", |
| 107 | DataTypeString(type)); |
| 108 | } |
| 109 | OP_REQUIRES_OK(ctx, output.status()); |
| 110 | |
| 111 | if (type == DT_INT32 || type == DT_INT64) { |
| 112 | // If input has dynamic dimension (value is -1), propagate the dynamic |
| 113 | // dimension to output using set-dimension-size. |
| 114 | ctx->set_dynamic_dimension_is_minus_one(true); |
| 115 | OP_REQUIRES_OK(ctx, ctx->ConstantInput(1, &limit)); |
| 116 | if (type == DT_INT32) { |
| 117 | if (limit.Get<int32>({}) == -1) { |
| 118 | output = xla::SetDimensionSize(output.ValueOrDie(), ctx->Input(1), 0); |
| 119 | } |
| 120 | } else { |
| 121 | if (limit.Get<int64>({}) == -1) { |
| 122 | output = xla::SetDimensionSize( |
| 123 | output.ValueOrDie(), |
| 124 | xla::ConvertElementType(ctx->Input(1), xla::S32), 0); |
| 125 | } |
| 126 | } |
| 127 | } |
| 128 | ctx->SetOutput(0, output.ValueOrDie()); |
| 129 | } |
nothing calls this directly
no test coverage detected