Skip to content

fix: unify LINSPACE alpha/beta dtype in the arange converter - #4626

Open
SrivastavaKshitij wants to merge 1 commit into
pytorch:mainfrom
SrivastavaKshitij:fix/arange-linspace-alpha-beta-dtype
Open

fix: unify LINSPACE alpha/beta dtype in the arange converter#4626
SrivastavaKshitij wants to merge 1 commit into
pytorch:mainfrom
SrivastavaKshitij:fix/arange-linspace-alpha-beta-dtype

Conversation

@SrivastavaKshitij

@SrivastavaKshitij SrivastavaKshitij commented Aug 28, 2026

Copy link
Copy Markdown

Description

Issue #4625 has two parts.

  1. Part 1 was solved on main by feat: Allow for users / kv cache to add aliased I/O for inplace operations #4251 and fix: arange converter #4456
  2. This PR fixes the second part which relates to dtype mismatch

arange resolves a common value_dtype for the sequence and passes it to each
get_trt_tensor call. get_trt_tensor applies that dtype when it constructs a constant, and
returns a value that is already an ITensor unchanged:

    elif isinstance(input_val, TRTTensor):
        return input_val

So a dynamic start keeps the dtype of its incoming ITensor, while a literal step becomes a
constant of value_dtype. TensorRT asks that the LINSPACE alpha (input 1) and beta
(input 2) have the same type, and reports:

IFillLayer `alpha` and `beta` must have the same type.
`alpha` is of type Int32 but `beta` is of type Int64.

This routes every operand through a small helper that casts after get_trt_tensor, so the
resolved dtype holds for all of them. cast_trt_tensor returns the tensor unchanged when the
dtype already matches, so paths that were already consistent are unaffected.

Issues

Closes #4625

Type of change

  • Bug fix (non-breaking change which fixes an issue)

Testing

Three cases with a dynamic bound are covered: an int32 start, an int64 start, and both
bounds dynamic. Of those, only the int64 start works today, because the sequence dtype for
integer operands is already int64 and matches.

Checklist:

  • My code follows the style guidelines of this project (You can use the linters)
  • I have performed a self-review of my own code
  • I have commented my code, particularly in hard-to-understand areas and hacks
  • I have made corresponding changes to the documentation
  • I have added tests to verify my fix or my feature
  • New and existing unit tests pass locally with my changes
  • I have added the relevant labels to my PR in so that relevant reviewers are notified

@meta-cla meta-cla Bot added the cla signed label Aug 28, 2026
@github-actions github-actions Bot added component: tests Issues re: Tests component: conversion Issues re: Conversion stage component: core Issues re: The core compiler component: converters Issues re: Specific op converters component: api [Python] Issues re: Python API component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels Aug 28, 2026
@github-actions
github-actions Bot requested a review from zewenli98 August 28, 2026 02:20
`arange` resolves a common `value_dtype` for the sequence and passes it to each
`get_trt_tensor` call. `get_trt_tensor` applies that dtype when it constructs a constant,
and returns a value that is already an ITensor unchanged:

    elif isinstance(input_val, TRTTensor):
        return input_val

So a dynamic `start` keeps the dtype of its incoming ITensor, while a literal `step`
becomes a constant of `value_dtype`. TensorRT asks that the LINSPACE `alpha` (input 1) and
`beta` (input 2) have the same type, and reports:

    IFillLayer `alpha` and `beta` must have the same type.
    `alpha` is of type Int32 but `beta` is of type Int64.

This change routes every operand through a small helper that casts after
`get_trt_tensor`, so the resolved dtype holds for all of them. `cast_trt_tensor` returns
the tensor unchanged when the dtype already matches, so paths that were already consistent
are unaffected.

The tests cover three cases with a dynamic bound: an int32 `start`, an int64 `start`, and
both bounds dynamic. Of those, only the int64 `start` works today, because the sequence
dtype for integer operands is already int64 and matches. The other two pass with this
change, and the full arange converter suite stays green.
@SrivastavaKshitij
SrivastavaKshitij force-pushed the fix/arange-linspace-alpha-beta-dtype branch from e3eeb41 to 91b21b9 Compare August 28, 2026 02:48
@SrivastavaKshitij

Copy link
Copy Markdown
Author

@narendasan @apbose for viz

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signed component: api [Python] Issues re: Python API component: conversion Issues re: Conversion stage component: converters Issues re: Specific op converters component: core Issues re: The core compiler component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: tests Issues re: Tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

🐛 [Bug] arange passes a rank-1 tensor to LINSPACE's rank-0 alpha input; the API error is logged, not raised

1 participant