Skip to content

ENH: Replace unsigned long with uint64_t for non-zero Jacobian indices #1316

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Apr 16, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Common/Transforms/itkAdvancedTransform.h

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@N-Dekker Thanks a lot for the modification! Just a quick note: torch::kLong in PyTorch corresponds to int64_t, so if the data is declared as uint64_t, torch::from_blob(..., torch::kLong) will interpret the underlying memory as signed. In practice, values from NonZeroJacobianIndicesType will likely never exceed INT64_MAX, so it should be fine, but it's something to keep in mind for strict type compatibility.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @vboussot Would it be possible to use torch::kUInt64 instead? 🤔 As in torch::from_blob(..., torch::kUInt64)?

Copy link

@vboussot vboussot Apr 16, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Unfortunately, LibTorch doesn’t support unsigned integer types beyond kUInt8 https://pytorch.org/docs/stable/tensor_attributes.html#torch.dtype

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Strange... I did find kUInt64 in the LibTorch 2.6.0 distribution:

libtorch-win-shared-with-deps-debug-2.6.0+cpu\libtorch\include\torch\csrc\api\include\torch\types.h from line 54:

constexpr auto kUInt64 = at::kUInt64;
constexpr auto kFloat16 = at::kHalf;
constexpr auto kFloat32 = at::kFloat;
constexpr auto kFloat64 = at::kDouble;

/// Rust-style short dtypes.
constexpr auto kU8 = kUInt8;
constexpr auto kU16 = kUInt16;
constexpr auto kU32 = kUInt32;
constexpr auto kU64 = kUInt64;

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@N-Dekker I tested it because the type torch::kUInt64 does exist, but in practice it’s not fully supported. For example, operations like index_add_() fail with the following error:

index_add_(): Expected dtype int32/int64 for index but got: UInt64

It seems that torch::from_blob(..., torch::kUInt64) works for reading the data, but a cast to torch::kLong is still required to ensure compatibility with indexing operations. I think I’ll go with that approach, since I’d rather keep uint64_t in elastix for types like indices, where unsigned values make more sense.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Casting a uint64 pointer to an int64 pointer, and then using its values "as if they are signed int64" technically sounds like undefined behavior. Right? On the other hand, it should be safe to just copy those Jacobian indices (uint64) into a signed int64 buffer, and pass that to Torch. Is it really a large performance issue?

If it's important enough, maybe we can make a pull request for torch to improve their kUInt64 support 😃

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I found this pytorch/pytorch#58734 it's a fairly recent addition in LibTorch

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure when kUInt64 was added, I think it was with PyTorch 2.3 (April 24, 2024). For now, it's reasonable to use torch::from_blob(..., torch::kLong) with uint64_t in this context. Thank you for the change

Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,7 @@ class ITK_TEMPLATE_EXPORT AdvancedTransform : public Transform<TScalarType, NInp
* Using an itk::FixedArray instead of an std::vector gives a performance
* gain for the SpatialHessianType.
*/
using NonZeroJacobianIndicesType = std::vector<unsigned long>;
using NonZeroJacobianIndicesType = std::vector<uint64_t>;
using SpatialJacobianType = Matrix<ScalarType, OutputSpaceDimension, InputSpaceDimension>;
using JacobianOfSpatialJacobianType = std::vector<SpatialJacobianType>;
// \todo: think about the SpatialHessian type, should be a 3D native type
Expand Down
4 changes: 2 additions & 2 deletions Common/Transforms/itkRecursiveBSplineTransform.hxx
Original file line number Diff line number Diff line change
Expand Up @@ -566,8 +566,8 @@ RecursiveBSplineTransform<TScalar, NDimensions, VSplineOrder>::ComputeNonZeroJac
}

/** Call the recursive implementation. */
unsigned long currentIndex = totalOffsetToSupportIndex;
unsigned long * nzjiPointer = &nonZeroJacobianIndices[0];
unsigned long currentIndex = totalOffsetToSupportIndex;
uint64_t * nzjiPointer = &nonZeroJacobianIndices[0];
ImplementationType::ComputeNonZeroJacobianIndices(nzjiPointer, parametersPerDim, currentIndex, gridOffsetTable);

} // end ComputeNonZeroJacobianIndices()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -122,7 +122,7 @@ class ITK_TEMPLATE_EXPORT RecursiveBSplineTransformImplementation

/** ComputeNonZeroJacobianIndices recursive implementation. */
static void
ComputeNonZeroJacobianIndices(unsigned long *& nzji,
ComputeNonZeroJacobianIndices(uint64_t *& nzji,
const unsigned long parametersPerDim,
unsigned long currentIndex,
const OffsetValueType * const gridOffsetTable)
Expand Down Expand Up @@ -409,7 +409,7 @@ class ITK_TEMPLATE_EXPORT RecursiveBSplineTransformImplementation<OutputDimensio

/** ComputeNonZeroJacobianIndices recursive implementation. */
static void
ComputeNonZeroJacobianIndices(unsigned long *& nzji,
ComputeNonZeroJacobianIndices(uint64_t *& nzji,
const unsigned long parametersPerDim,
const unsigned long currentIndex,
const OffsetValueType * const itkNotUsed(gridOffsetTable))
Expand Down
Loading