Features/714 trace

[lenablind]

Description

Implementation of ht.trace, which computes the sum along the diagonal, analogous to np.trace.

Used torch functions:

• torch.trace
• torch.diagonal
• torch.sum

Issue/s resolved: #714

Strategy (Algorithm)

CASE 2D Input
Returns: Scalar

• If offset would result into an empty array, directly create a 0d-tensor with value 0
• Otherwise, adapt the offset to the distribution,
• Defer the resulting sub-arrays, upholding the affected diagonal entries then
• Calculate the trace of each sub-array and
• Sum up all partial results by means of an MPI_Allreduce operation to obtain the final result

CASE > 2D Input
Returns: DNDArray
(Hint: as torch.trace is only implemented for 2-dimensional tensors, I decided to do a workaround by first extracting the diagonal entries via torch.diagonal followed by a summation along the last axis of the resulting array)

I differentiated between two cases: the split axis being or not being within the "trace axes" (axis1, axis2), as this leads to the required diagonal entries of the last axis being distributed on varying processes or not.

CASE: split axis NOT in (axis1, axis2)

• directly extract the diagonal of the data chunk
• compute the sum along the last axis via torch.sum as explained above
• set the partial results back together along the correct axis ("gather_axis")

CASE: split axis IN (axis1, axis2)

• adapt the offset to the distribution*
• (*: if it would not result into an empty array, otherwise torch.diagonal provides such even with unmodified offset)
• extract the diagonal entries
• adapt the shape of this array to the shape of result
• sum up all elements by means of MPI_Allgather

If out is not None, provide the correct split configuration to store the result in the given variable.

Changes proposed:

• Implementation of ht.trace
• Implementation of corresponding test

Type of change

• New feature (non-breaking change which adds functionality)

Due Diligence

• All split configurations tested
• Multiple dtypes tested in relevant functions
• Documentation updated (if needed)
• Updated changelog.md under the title "Pending Additions"

Does this change modify the behaviour of other functions? If so, which?

no

[coquelin77]

rerun tests

[coquelin77]

to clarify, you will return a split object even if the input is split=None?

[lenablind]

No, this won't happen. If a.split=None, the first part of the (and) if condition is already false and the else case is always executed. Thus the result will also be not distributed.

[coquelin77]

rerun tests