Skip to content

Commit 6ff2674

Browse files
committed
Updated concatenate access
1 parent 72dacee commit 6ff2674

File tree

1 file changed

+30
-6
lines changed

1 file changed

+30
-6
lines changed

include/xtensor/xbuilder.hpp

Lines changed: 30 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -497,20 +497,44 @@ namespace xt
497497
template <class It>
498498
inline value_type access(const tuple_type& t, size_type axis, It first, It last) const
499499
{
500-
xindex index(first, last);
501-
auto match = [&index, axis](auto& arr)
500+
// trim off extra indices if provided to match behavior of containers
501+
auto dim_offset = std::distance(first, last) - std::get<0>(t).dimension();
502+
size_t axis_dim = *(first + axis + dim_offset);
503+
auto match = [&](auto& arr)
502504
{
503-
if (index[axis] >= arr.shape()[axis])
505+
if (axis_dim >= arr.shape()[axis])
504506
{
505-
index[axis] -= arr.shape()[axis];
507+
axis_dim -= arr.shape()[axis];
506508
return false;
507509
}
508510
return true;
509511
};
510512

511-
auto get = [&index](auto& arr)
513+
auto get = [&](auto& arr)
512514
{
513-
return arr[index];
515+
size_t offset = 0;
516+
const size_t end = arr.dimension();
517+
for (size_t i = 0; i < end; i++)
518+
{
519+
const auto& shape = arr.shape();
520+
const size_t stride = std::accumulate(
521+
shape.begin() + i + 1,
522+
shape.end(),
523+
1,
524+
std::multiplies<size_t>()
525+
);
526+
if (i == axis)
527+
{
528+
offset += axis_dim * stride;
529+
}
530+
else
531+
{
532+
const auto len = (*(first + i + dim_offset));
533+
offset += len * stride;
534+
}
535+
}
536+
const auto element = arr.begin() + offset;
537+
return *element;
514538
};
515539

516540
size_type i = 0;

0 commit comments

Comments
 (0)