Skip to content

Refactor NumberOfElements dtype dispatch - #4064

Merged
zcbenz merged 1 commit into
ml-explore:mainfrom
PhysicistJohn:dispatch/number-of-elements
Aug 8, 2026
Merged

Refactor NumberOfElements dtype dispatch#4064
zcbenz merged 1 commit into
ml-explore:mainfrom
PhysicistJohn:dispatch/number-of-elements

Conversation

@PhysicistJohn

Copy link
Copy Markdown
Contributor

Replace the exhaustive dtype switch in NumberOfElements::eval with the existing dispatch_all_types utility.

This preserves the supported dtypes and existing conversion behavior while centralizing the type mapping. It follows the small-PR dtype-dispatch cleanup suggested in #3341.

Testing

  • Metal-enabled Python wheel build
  • Serial native C++ test suite
  • Shapeless compiled mean on CPU and GPU
  • pre-commit formatting checks

Comment thread mlx/backend/common/common.cpp Outdated
*out.data<complex64_t>() = static_cast<complex64_t>(numel);
break;
}
dispatch_all_types(out.dtype(), [&]<typename T>(type_identity<T>) {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Thanks for doing the refactoring, though I'm not sure this is correct, it does not match other dispatch_all_types calls.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Good catch. I updated this to match the existing auto type_tag / MLX_GET_TYPE(type_tag) pattern, and made the same change in #4065 and #4066. The focused builds and tests still pass.

@PhysicistJohn
PhysicistJohn force-pushed the dispatch/number-of-elements branch from 90f515b to c278498 Compare August 8, 2026 06:05
@zcbenz
zcbenz merged commit 361883c into ml-explore:main Aug 8, 2026
28 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants