Method select

bmf/hmp/src/tensor.cpp:243–266

}

Tensor Tensor::select(int64_t dim, int64_t index) const {
    dim = wrap_size(dim, this->dim());
    HMP_REQUIRE(dim < this->dim(), "select : dim({}) is out of range {}", dim,
                this->dim());

    index = wrap_size(index, size(dim));
    HMP_REQUIRE(index < size(dim),
                "select: index {} is out of range, need less than {}", index,
                size(dim));

    auto out = slice(dim, index, index + 1);

    // squeeze
    auto newShape = out.shape();
    auto newStrides = out.strides();
    HMP_REQUIRE(newShape[dim] == 1, "select: internal error");
    newShape.erase(newShape.begin() + dim);
    newStrides.erase(newStrides.begin() + dim);

    out.tensorInfo()->setSizesAndStrides(newShape, newStrides,
                                         out.tensorInfo()->bufferOffset());

    return out;
}

// factory functions
Tensor empty(const SizeArray &shape, const TensorOptions &options) {