| 49 | } |
| 50 | |
| 51 | void Compile(XlaOpKernelContext* ctx) override { |
| 52 | const TensorShape input_shape = ctx->InputShape(0); |
| 53 | const TensorShape tensor_axis_shape = ctx->InputShape(1); |
| 54 | |
| 55 | OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(tensor_axis_shape), |
| 56 | errors::InvalidArgument("ScanOp: axis must be a scalar, not ", |
| 57 | tensor_axis_shape.DebugString())); |
| 58 | |
| 59 | int64 axis; |
| 60 | OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntScalar(1, &axis)); |
| 61 | if (axis < 0) { |
| 62 | axis += input_shape.dims(); |
| 63 | } |
| 64 | OP_REQUIRES( |
| 65 | ctx, FastBoundsCheck(axis, input_shape.dims()), |
| 66 | errors::InvalidArgument("ScanOp: Expected scan axis in the range [", |
| 67 | -input_shape.dims(), ", ", input_shape.dims(), |
| 68 | "), but got ", axis)); |
| 69 | |
| 70 | DataType dtype = XlaHelpers::SumAccumulationType(ctx->input_type(0)); |
| 71 | |
| 72 | if (input_shape.num_elements() == 0) { |
| 73 | // Exit early if there is nothing to compute. |
| 74 | ctx->SetOutput(0, ctx->Input(0)); |
| 75 | return; |
| 76 | } |
| 77 | |
| 78 | xla::XlaBuilder* builder = ctx->builder(); |
| 79 | |
| 80 | std::vector<int64> window_strides(input_shape.dims(), 1); |
| 81 | std::vector<int64> window_dims(input_shape.dims(), 1); |
| 82 | window_dims[axis] = input_shape.dim_size(axis); |
| 83 | |
| 84 | std::vector<std::pair<int64, int64>> padding(input_shape.dims(), {0, 0}); |
| 85 | padding[axis].first = input_shape.dim_size(axis) - 1; |
| 86 | // In exclusive mode, add an extra padding element so there is a complete |
| 87 | // window of padding before the data starts. |
| 88 | if (exclusive_) { |
| 89 | ++padding[axis].first; |
| 90 | } |
| 91 | if (reverse_) { |
| 92 | std::swap(padding[axis].first, padding[axis].second); |
| 93 | } |
| 94 | |
| 95 | xla::XlaOp init; |
| 96 | const xla::XlaComputation* reducer; |
| 97 | if (sum_) { |
| 98 | init = XlaHelpers::Zero(builder, dtype); |
| 99 | reducer = ctx->GetOrCreateAdd(dtype); |
| 100 | } else { |
| 101 | init = XlaHelpers::One(builder, dtype); |
| 102 | reducer = ctx->GetOrCreateMul(dtype); |
| 103 | } |
| 104 | auto output = xla::ReduceWindowWithGeneralPadding( |
| 105 | XlaHelpers::ConvertElementType(ctx->Input(0), dtype), init, *reducer, |
| 106 | window_dims, window_strides, |
| 107 | /*base_dilations=*/{}, /*window_dilations=*/{}, padding); |
| 108 | output = XlaHelpers::ConvertElementType(output, ctx->input_type(0)); |
nothing calls this directly
no test coverage detected