Can remap_reduce_types be removed

[zcbenz]

The reduce kernels use remap_reduce_types to deduce the dtypes of input and output:

mlx/mlx/backend/metal/reduce.cpp

Lines 309 to 310 in 4c1dfa5

	// Set the kernel
	auto [in_type, out_type] = remap_reduce_types(in, op_name);

Which I think is redundant considering the op has already set dtype for them:

mlx/mlx/ops.cpp

Lines 1855 to 1869 in 4c1dfa5

	Dtype out_type = a.dtype();
	if (issubdtype(a.dtype(), signedinteger)) {
	out_type = a.dtype().size() <= 4 ? int32 : int64;
	} else if (issubdtype(a.dtype(), unsignedinteger)) {
	out_type = a.dtype().size() <= 4 ? uint32 : uint64;
	} else if (a.dtype() == bool_) {
	out_type = int32;
	}
	auto out = (is_noop)
	? astype(a, out_type, s)
	: array(
	std::move(out_shape),
	out_type,
	std::make_shared<Reduce>(to_stream(s), Reduce::Sum, sorted_axes),
	{a});

Why use remap_reduce_types instead of in/out.dtype()?

[angeloskath]

Oh that's an interesting one. I wouldn't say it is a hack per-se but we reduce the number of kernels (which reduces the binary size) by reusing the reduction kernels based on the size of the dtype. Basically that means that we use the int8 type for both uint8 and int8 and so on and so forth.

[zcbenz]

Thanks!