Skip to content

Preserve tensor dtype when the left operand is an R scalar#1471

Draft
TroyHernandez wants to merge 2 commits into
mlverse:mainfrom
cornball-ai:fix/scalar-lhs-type-promotion
Draft

Preserve tensor dtype when the left operand is an R scalar#1471
TroyHernandez wants to merge 2 commits into
mlverse:mainfrom
cornball-ai:fix/scalar-lhs-type-promotion

Conversation

@TroyHernandez

Copy link
Copy Markdown
Contributor

Problem

Arithmetic with an R scalar on the left silently promotes reduced-precision and integer tensors:

t <- torch_ones(3, dtype = torch_bfloat16())
(t - 1)$dtype   # BFloat16, as expected
(1 - t)$dtype   # Float  <- silently upcast
(1 - torch_ones(3, dtype = torch_long()))$dtype   # Float (PyTorch returns Long)

In Python, 1 - t keeps t's dtype because Python scalars participate weakly in type promotion.

This surfaced as a hard failure in a bfloat16 attention stack: an additive mask built with 1 - attention_mask came out float32, and CUDA's fused scaled_dot_product_attention rejects a bias whose dtype doesn't match the query ("invalid dtype for bias - should match query's dtype").

Cause

The operator methods wrap a non-tensor lhs with torch_tensor(e1, ...), which yields a dimensioned (length-1, 1-d) tensor of the default dtype. libtorch lets dimensioned tensors dictate promotion, so float32 wins over bfloat16/int64. The rhs case never hits this path: torch_sub(tensor, 1) takes the Tensor-Scalar overload, which promotes weakly.

Fix

Wrap length-1 numeric/logical lhs operands with torch_scalar_tensor() instead. 0-dim tensors promote weakly, exactly like Python scalars, so:

  • 1 - bf16 stays BFloat16, 2 * half stays Half
  • 1L + long stays Long, 5L %% long stays Long
  • 1L / long still promotes to the default float dtype (division semantics unchanged)
  • c(1, 2, 3) - bf16 still promotes to Float (dimensioned vectors keep their old behavior)
  • NA, NaN, logical, and integer scalars behave as before (verified against torch_tensor() conversions)

Applied to the seven arithmetic ops (+, -, *, /, ^, %%, %/%). Comparison operators share the wrapping pattern but return bool, so they're left untouched.

The regression test fails 4 assertions before the patch and passes after; the rest of test-operators.R is unchanged (149 passing before and after, CPU and CUDA checked).

`1 - t` converted the scalar with torch_tensor(), producing a
dimensioned float32 tensor that dictates type promotion: bfloat16 and
float16 tensors silently upcast to float32, and `1 - long_tensor`
returns float. Scalars on the rhs already hit the Tensor-Scalar
overloads with weak promotion, so `t - 1` kept the dtype while
`1 - t` did not. Convert length-1 numeric/logical lhs operands with
torch_scalar_tensor() instead: 0-dim tensors promote weakly, matching
Python's scalar semantics.
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.

1 participant