Fix vmap + floor_divide: preserve integer dtype

[robert-johansson]

Summary

Divide::vmap calls divide(a, b) which promotes integers to float via at_least_float. But floor_divide creates a Divide primitive with integer dtype for integer division. Under vmap, the integer dtype is lost.

Before:

import mlx.core as mx
f = lambda s: mx.floor_divide(s, mx.array(5, mx.int32))
vf = mx.vmap(f)
result = vf(mx.arange(25, dtype=mx.int32))
# result.dtype == float32, values are [0.0, 0.2, 0.4, ...] (wrong)

After:

result = vf(mx.arange(25, dtype=mx.int32))
# result.dtype == int32, values are [0, 0, 0, 0, 0, 1, 1, ...] (correct)

Fix

In Divide::vmap, check if inputs are integer type and construct the Divide primitive directly (preserving dtype) instead of calling the high-level divide() which promotes to float.

Test plan

• New test case test vmap floor_divide integer with 3 sub-tests:
  • floor_divide preserves int32 dtype under vmap
  • remainder preserves int32 dtype under vmap
  • floor_divide + remainder reconstructs original values under vmap
• All existing vmap tests pass (38 assertions)

[zcbenz]

I think this would make vmap + divide output wrong dtype for integers? Adding a flag in Divide might be a better solution?