apple / apple/coremltools

`torch.arange` ignores `dtype` when used statically

Open
#1,852 1 comment 1 reaction 0 assignees View on GitHub
bug PyTorch (traced) triaged
Dominant language
Python
Stars
5.4k
Forks
850
Avg merge
4d 5h
Merged PRs (30d)
10

Description

## 🐞Describing the bug
`torch.arange(value, dtype=torch.float)` yields an `int` tensor instead of `float`, as demonstrated in the code snippet below, if `value` is an integer. This produces conversion errors in several transformers models where sinusoidal positions can be created like this: https://github.com/huggingface/transformers/blob/b61d5b47f640308068139561f673765b2af39874/src/transformers/models/gptj/modeling_gptj.py#L60

## Stack Trace
Reveal

```
Traceback (most recent call last):
File "/opt/homebrew/Caskroom/miniforge/base/envs/sdcoreml/lib/python3.9/runpy.py", line 197, in _run_module_as_main
return _run_code(code, main_globals, None,
File "/opt/homebrew/Caskroom/miniforge/base/envs/sdcoreml/lib/python3.9/runpy.py", line 87, in _run_code
exec(code, run_globals)
File "/Users/pedro/.vscode/extensions/ms-python.python-2023.6.1/pythonFiles/lib/python/debugpy/adapter/../../debugpy/launcher/../../debugpy/__main__.py", line 39, in
cli.main()
File "/Users/pedro/.vscode/extensions/ms-python.python-2023.6.1/pythonFiles/lib/python/debugpy/adapter/../../debugpy/launcher/../../debugpy/../debugpy/server/cli.py", line 430, in main
run()
File "/Users/pedro/.vscode/extensions/ms-python.python-2023.6.1/pythonFiles/lib/python/debugpy/adapter/../../debugpy/launcher/../../debugpy/../debugpy/server/cli.py", line 284, in run_file
runpy.run_path(target, run_name="__main__")
File "/Users/pedro/.vscode/extensions/ms-python.python-2023.6.1/pythonFiles/lib/python/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 321, in run_path
return _run_module_code(code, init_globals, run_name,
File "/Users/pedro/.vscode/extensions/ms-python.python-2023.6.1/pythonFiles/lib/python/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 135, in _run_module_code
_run_code(code, mod_globals, init_globals,
File "/Users/pedro/.vscode/extensions/ms-python.python-2023.6.1/pythonFiles/lib/python/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 124, in _run_code
exec(code, run_globals)
File "/Users/pedro/code/hf/apple/exporters/einsum_test.py", line 20, in
mlmodel = ct.convert(
File "/Users/pedro/code/hf/apple/coremltools/coremltools/converters/_converters_entry.py", line 492, in convert
mlmodel = mil_convert(
File "/Users/pedro/code/hf/apple/coremltools/coremltools/converters/mil/converter.py", line 188, in mil_convert
return _mil_convert(model, convert_from, convert_to, ConverterRegistry, MLModel, compute_units, **kwargs)
File "/Users/pedro/code/hf/apple/coremltools/coremltools/converters/mil/converter.py", line 212, in _mil_convert
proto, mil_program = mil_convert_to_proto(
File "/Users/pedro/code/hf/apple/coremltools/coremltools/converters/mil/converter.py", line 285, in mil_convert_to_proto
prog = frontend_converter(model, **kwargs)
File "/Users/pedro/code/hf/apple/coremltools/coremltools/converters/mil/converter.py", line 108, in __call__
return load(*args, **kwargs)
File "/Users/pedro/code/hf/apple/coremltools/coremltools/converters/mil/frontend/torch/load.py", line 63, in load
return _perform_torch_convert(converter, debug)
File "/Users/pedro/code/hf/apple/coremltools/coremltools/converters/mil/frontend/torch/load.py", line 102, in _perform_torch_convert
prog = converter.convert()
File "/Users/pedro/code/hf/apple/coremltools/coremltools/converters/mil/frontend/torch/converter.py", line 284, in convert
convert_nodes(self.context, self.graph)
File "/Users/pedro/code/hf/apple/coremltools/coremltools/converters/mil/frontend/torch/ops.py", line 88, in convert_nodes
add_op(context, node)
File "/Users/pedro/code/hf/apple/coremltools/coremltools/converters/mil/frontend/torch/ops.py", line 1122, in einsum
x = build_einsum_mil(a, b, equation, node.name)
File "/Users/pedro/code/hf/apple/coremltools/coremltools/converters/mil/frontend/_utils.py", line 171, in build_einsum_mil
x = solve_generic_einsum(parsed_vectors, a_var, b_var, name)
File "/Users/pedro/code/hf/apple/coremltools/coremltools/converters/mil/frontend/_utils.py", line 390, in solve_generic_einsum
ab = mb.matmul(x=a, y=b)
File "/Users/pedro/code/hf/apple/coremltools/coremltools/converters/mil/mil/ops/registry.py", line 182, in add_op
return cls._add_op(op_cls_to_add, **kwargs)
File "/Users/pedro/code/hf/apple/coremltools/coremltools/converters/mil/mil/builder.py", line 166, in _add_op
new_op = op_cls(**kwargs)
File "/Users/pedro/code/hf/apple/coremltools/coremltools/converters/mil/mil/operation.py", line 187, in __init__
self._validate_and_set_inputs(input_kv)
File "/Users/pedro/code/hf/apple/coremltools/coremltools/converters/mil/mil/operation.py", line 496, in _validate_and_set_inputs
self.input_spec.validate_inputs(self.name, self.op_type, input_kvs)
File "/Users/pedro/code/hf/apple/coremltools/coremltools/converters/mil/mil/input_type.py", line 137, in validate_inputs
raise ValueError(msg)
ValueError: In op, of type matmul, named matmul_0, the named input `y` must have the same data type as the named input `x`. However, y has dtype fp32 whereas x has dtype int32.
```

## To Reproduce
```Python
import torch
import torch.nn as nn
import coremltools as ct

class MinimalRepro(nn.Module):
def forward(self, x):
# Outer product
return torch.einsum("i,j->ij", torch.arange(5, dtype=torch.float), x)

model = MinimalRepro()

x = torch.randn(3)
assert model(x).shape == (5, 3)

jitted_model = torch.jit.trace(model, x)
jitted_model.eval()
x_j = jitted_model(x)
assert (x_j - model(x)).abs().max() < 1e-4

mlmodel = ct.convert(
jitted_model,
inputs=[ct.TensorType(name="x", shape=x.shape)],
)
```

Observe that the previous example works if any of the following modifications are performed:

- Use a float value as an argument to `arange`:

` return torch.einsum("i,j->ij", torch.arange(5.0, dtype=torch.float), x)`

- Explicitly convert to `float` anyway:

` return torch.einsum("i,j->ij", torch.arange(5, dtype=torch.float).float(), x)`

In addition, the following also works:

```Python
import torch
import torch.nn as nn
import coremltools as ct

class MinimalRepro(nn.Module):
def __init__(self):
super().__init__()
self.y = torch.randn(3)

def forward(self, x):
# Outer product
return torch.einsum("i,j->ij", x, self.y)

model = MinimalRepro()

x = torch.arange(5, dtype=float)
assert model(x).shape == (5, 3)

jitted_model = torch.jit.trace(model, x)
jitted_model.eval()
x_j = jitted_model(x)
assert (x_j - model(x)).abs().max() < 1e-4

mlmodel = ct.convert(
jitted_model,
inputs=[ct.TensorType(name="x", shape=x.shape)],
)
```

## System environment (please complete the following information):
- coremltools version: 6.3.0
- OS (e.g. MacOS version or Linux type): macOS Ventura 13.3, 13.1
- Any other relevant version information (e.g. PyTorch or TensorFlow version): PyTorch 1.13.1

## Additional context
In the presence of the bug, registering this custom op and setting a breakpoint confirms that the type of `a` is integer (and therefore `dtype` was ignored):

```Python
from coremltools.converters.mil.frontend.torch.torch_op_registry import _TORCH_OPS_REGISTRY, register_torch_op

del _TORCH_OPS_REGISTRY["einsum"]

@register_torch_op
def einsum(context, node):
from coremltools.converters.mil import Builder as mb
from coremltools.converters.mil.frontend._utils import build_einsum_mil
from coremltools.converters.mil.mil import types

a = context[node.inputs[1]][0]
b = context[node.inputs[1]][1]
equation = context[node.inputs[0]].val
equation = "".join(equation.split(" "))
if equation == "i,j->ij" and types.is_int(a.dtype):
a = mb.cast(x=a, dtype="fp32")
x = build_einsum_mil(a, b, equation, node.name)

context.add(x)
```

Contributor guide

Open the contributing guide

Research direction

Start with the Torch frontend conversion path in coremltools/converters/mil/frontend/torch/ops.py, especially the arange handling, and compare it with the torch op registry entry points. Reproduce the issue with the MinimalRepro example and confirm that a static torch.arange call preserves dtype=float and allows the einsum conversion to complete without the dtype mismatch.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, tooling
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.