Replace uint8 Masks with Bool Tensors

[tmattio]

Background

Nx recently added support for extended data types including boolean arrays. However, comparison and conditional operations still use uint8 tensors instead of proper boolean tensors.

Objective

Update comparison and conditional operations to use boolean tensors:

• Change comparison ops (op_cmplt, op_cmpne, etc.) to return (bool, Dtype.bool_elt) t
• Update op_where to accept boolean condition tensors
• Migrate all mask/condition usage in frontend operations

Design

Current signatures:

val op_cmplt : ('a, 'b) t -> ('a, 'b) t -> (int, Dtype.uint8_elt) t
val op_where : (int, Dtype.uint8_elt) t -> ('a, 'b) t -> ('a, 'b) t -> ('a, 'b) t

New signatures:

val op_cmplt : ('a, 'b) t -> ('a, 'b) t -> (bool, Dtype.bool_elt) t
val op_where : (bool, Dtype.bool_elt) t -> ('a, 'b) t -> ('a, 'b) t -> ('a, 'b) t

Implementation Areas

1. Backend Interface (nx/lib/core/backend_intf.ml)
2. Native Backend (nx/lib/native/)
3. C Backend (nx/lib/c/)
4. Metal Backend (nx/lib/metal/)
5. Frontend (nx/lib/core/frontend.ml)

Testing

• Ensure all existing tests pass with bool tensors

[furmak331]

Clean issue with clear requirements. The bool tensor migration makes total sense from a type safety perspective.

I'd like to take this on. I will try to implement in every backend interface systematically and check if they pass the tests.

[nirnayroy]

@furmak331 , Are you still working on this?
If not, do you mind if I take this up?

cc: @tmattio

[tmattio]

You can take it on @nirnayroy, although note that it wasn't initially an issue for Outreachy applicant, and might not be a trivial one.

[nirnayroy]

Got it!