Method Compile

tensorflow/compiler/tf2xla/kernels/gather_op.cc:164–238

  }

  void Compile(XlaOpKernelContext* context) override {
    auto input = context->Input(0);
    auto input_shape = context->InputShape(0);
    auto indices = context->Input(1);
    auto indices_shape = context->InputShape(1);

    absl::optional<int64> axis;
    if (context->num_inputs() == 3) {
      const TensorShape axis_shape = context->InputShape(2);
      OP_REQUIRES(context, TensorShapeUtils::IsScalar(axis_shape),
                  errors::InvalidArgument("axis must be scalar"));
      DataType axis_type = input_type(2);
      OP_REQUIRES(context, axis_type == DT_INT32 || axis_type == DT_INT64,
                  errors::InvalidArgument("axis must be int32 or int64"));

      int64 axis_input;
      OP_REQUIRES_OK(context,
                     context->ConstantInputAsIntScalar(2, &axis_input));

      const auto params_dims = input_shape.dims();
      OP_REQUIRES(context,
                  -params_dims <= axis_input && axis_input < params_dims,
                  errors::InvalidArgument("Expected axis in the range [",
                                          -params_dims, ", ", params_dims,
                                          "), but got ", axis_input));
      if (axis_input < 0) {
        axis_input += params_dims;
      }
      axis = axis_input;
    }

    if (batch_dims_ != 0) {
      if (batch_dims_ < 0) {
        batch_dims_ = indices_shape.dims() + batch_dims_;
      }

      axis = axis.value_or(batch_dims_);

      OP_REQUIRES(context,
                  batch_dims_ >= -indices_shape.dims() &&
                      batch_dims_ < indices_shape.dims(),
                  errors::InvalidArgument("Expected batch_dims in the range [",
                                          -indices_shape.dims(), ", ",
                                          indices_shape.dims(), "), but got ",
                                          batch_dims_));

      OP_REQUIRES(context, batch_dims_ < input_shape.dims(),
                  errors::InvalidArgument("batch_dims (", batch_dims_,
                                          ") must be less than rank(input) (",
                                          input_shape.dims(), ")."));

      OP_REQUIRES(context, *axis >= batch_dims_,
                  errors::InvalidArgument("batch_dims (", batch_dims_,
                                          ") must be less than or equal to ",
                                          "axis (", *axis, ")."));
    }

    axis = axis.value_or(0);