Add torch op coverage for LLM attention mask construction

[john-rocky]

Summary

Adds the small set of torch ops that HuggingFace attention-mask construction emits via torch.export but coremltools didn't yet handle. Discovered while converting google/gemma-4-E2B-it, but the same gaps affect any decoder LLM that builds masks the way HuggingFace's eager attention path does.

All five gaps are independent fixes for existing asymmetries in the registry (e.g. bitwise_and exists but bitwise_or did not; new_zeros exists but new_ones did not). None of them change behavior of any op that already worked.

Reproduction (before this PR)

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import coremltools as ct

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-4-E2B-it",
    torch_dtype=torch.float32,
    attn_implementation="eager",
).eval()
tok = AutoTokenizer.from_pretrained("google/gemma-4-E2B-it")
enc = tok("hello", return_tensors="pt", padding="max_length", max_length=32)

class Wrapper(torch.nn.Module):
    def __init__(self, m):
        super().__init__()
        self.m = m
    def forward(self, ids, attn):
        return self.m(input_ids=ids, attention_mask=attn, use_cache=False).logits

with torch.no_grad():
    ep = torch.export.export(
        Wrapper(model),
        args=(enc["input_ids"], enc["attention_mask"]),
        strict=False,
    ).run_decompositions({})

ct.convert(ep, convert_to="mlprogram", minimum_deployment_target=ct.target.iOS18)
# NotImplementedError: Unsupported fx node or_1, kind __or__

After fixing __or__ the converter then fails on __and__, then new_ones, then where.scalarother, then bitwise_and with mixed bool+float inputs. This PR addresses all of them in one place.

Changes

1. sanitize_op_kind — strip the __name__ wrapper after the namespace and overload suffix have been removed. Previously aten::__or__.Tensor sanitized to __or__ instead of or, so the registry lookup missed even when an or handler existed. The original logic only stripped __...__ if the whole op kind matched, which is only the case for the legacy TorchScript form.

2. Register bitwise_or / bitwise_xor with or / xor aliases for the post-sanitize form. Both reuse the logical_* MIL primitives, mirroring the existing bitwise_and registration. The asymmetry was clearly accidental.

3. Relax bitwise_and / bitwise_or / bitwise_xor to accept mixed-bool inputs. Previously they required all inputs to be bool; now if at least one input is bool we cast both to bool and lower to the corresponding logical_* op. Pure non-bool inputs are still rejected with the same error so genuine integer bitwise math is unchanged. This unblocks Gemma-style mask combination where a bool causal mask meets a float padding mask.

4. Register aten::new_ones mirroring the existing new_zeros, using _make_fill_op so float-typed shape inputs from torch.export are coerced to int32 (which is also a latent issue in new_zeros that this PR avoids by using the safer helper).

5. Add where.ScalarOther as an alias on the existing where handler, which already does dtype promotion and broadcasting. The where.self and where.scalarself overloads were already covered; where.scalarother was the missing third.

Test plan

• New unit tests in test_internal_graph.py covering the dunder-after-namespace case for sanitize_op_kind:
  • aten::__or__.Tensor → or
  • aten::__and__.Tensor → and
  • aten::__xor__.Tensor → xor
  • Plus the existing legacy forms (__add__ → add, aten::add.Tensor → add, etc.) to guard against regressions
• New op-level tests in test_torch_ops.py:
  • TestNewOnes::test_new_ones_static
  • TestBitwiseOr::test_bitwise_or and test_or_operator (the latter exercises the tensor | tensor form)
  • TestBitwiseXor::test_bitwise_xor
• Existing TestBitwiseAnd and all TestLogical* tests still pass (the relaxed bool check is a strict superset of the previous behavior)

End-to-end validation on google/gemma-4-E2B-it (compute_precision=FLOAT32):

ref last argmax: 7001 -> ' France'
cml last argmax: 7001 -> ' France'
top-5 overlap: 5/5
per-position argmax agreement: 100.0% (5/5)
max abs diff: 5.30e-02
VERDICT: Core ML model is functionally correct.

(fp16 correctness for the same model is addressed by a separate PR: #2669 — no functional dependency between the two PRs, but both are needed to convert Gemma 4 cleanly.)

[TobyRoseman]

CI Run: https://gitlab.com/coremltools1/coremltools/-/pipelines/2450002299

[TobyRoseman]

Thanks for your contribution @john-rocky.