bug: torch 2.x with TorchTensor

PyTorch 2.0 adds extensive compilation support to the framework; we should test this with our TorchTensor implementations and make sure that

it works
we don't somehow negate the performance gains

Hi! I'd like to work on this issue

Hi! I'd like to work on this issue

Sure thing! I just discovered that I cannot formally assign it to you bc I am on mobile, but consider it yours!

@h4shk4t

Hi @h4shk4t , have you had time to look into this?

Hey yes I haven't gone through it properly as my laptop is gone for repair. I will get it back it in 2 days. I can try to work on it regardless if it needs to be completed urgently

Issue back open for whoever wants to take it ^^

So I gave this a try, and sadly it doesn't work:

import torch
from docarray.typing import TorchTensor
from pydantic import parse_obj_as


def fn(a, b):
    return a @ b


# native
a, b = torch.rand(1000, 1000), torch.rand(1000, 1000)
fn_c = torch.compile(fn)
fn_c(a, b)
print('native successful')

# docarray
a, b = parse_obj_as(TorchTensor[1000, 1000], torch.rand(1000, 1000)), parse_obj_as(TorchTensor[1000, 1000], torch.rand(1000, 1000))
fn_c = torch.compile(fn)
fn_c(a, b)
print('docarray successful')

native successful
Traceback (most recent call last):
  File "/home/johannes/Documents/jina/docarrayv2/venv/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py", line 324, in _compile
    out_code = transform_code_object(code, transform)
  File "/home/johannes/Documents/jina/docarrayv2/venv/lib/python3.8/site-packages/torch/_dynamo/bytecode_transformation.py", line 445, in transform_code_object
    transformations(instructions, code_options)
  File "/home/johannes/Documents/jina/docarrayv2/venv/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py", line 299, in transform
    tracer = InstructionTranslator(
  File "/home/johannes/Documents/jina/docarrayv2/venv/lib/python3.8/site-packages/torch/_dynamo/symbolic_convert.py", line 1670, in __init__
    self.symbolic_locals = collections.OrderedDict(
  File "/home/johannes/Documents/jina/docarrayv2/venv/lib/python3.8/site-packages/torch/_dynamo/symbolic_convert.py", line 1673, in <genexpr>
    VariableBuilder(
  File "/home/johannes/Documents/jina/docarrayv2/venv/lib/python3.8/site-packages/torch/_dynamo/variables/builder.py", line 172, in __call__
    return self._wrap(value).clone(**self.options())
  File "/home/johannes/Documents/jina/docarrayv2/venv/lib/python3.8/site-packages/torch/_dynamo/variables/builder.py", line 453, in _wrap
    istype(value, (type, types.FunctionType))
  File "/home/johannes/Documents/jina/docarrayv2/venv/lib/python3.8/site-packages/torch/_dynamo/utils.py", line 346, in istype
    return type(obj) in allowed_types
  File "/home/johannes/Documents/jina/docarrayv2/docarray/typing/tensor/abstract_tensor.py", line 101, in __eq__
    if cls._equals_special_case(other):
  File "/home/johannes/Documents/jina/docarrayv2/docarray/typing/tensor/abstract_tensor.py", line 62, in _equals_special_case
    is_tensor = is_type and AbstractTensor in other.mro()
TypeError: descriptor 'mro' of 'type' object needs an argument

Set torch._dynamo.config.verbose=True for more information


You can suppress this exception and fall back to eager by setting:
    torch._dynamo.config.suppress_errors = True


The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/johannes/.config/JetBrains/PyCharmCE2023.2/scratches/scratch_182.py", line 18, in <module>
    fn_c(a, b)
  File "/home/johannes/Documents/jina/docarrayv2/venv/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py", line 209, in _fn
    return fn(*args, **kwargs)
  File "/home/johannes/Documents/jina/docarrayv2/venv/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py", line 337, in catch_errors
    return callback(frame, cache_size, hooks)
  File "/home/johannes/Documents/jina/docarrayv2/venv/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py", line 404, in _convert_frame
    result = inner_convert(frame, cache_size, hooks)
  File "/home/johannes/Documents/jina/docarrayv2/venv/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py", line 104, in _fn
    return fn(*args, **kwargs)
  File "/home/johannes/Documents/jina/docarrayv2/venv/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py", line 262, in _convert_frame_assert
    return _compile(
  File "/home/johannes/Documents/jina/docarrayv2/venv/lib/python3.8/site-packages/torch/_dynamo/utils.py", line 163, in time_wrapper
    r = func(*args, **kwargs)
  File "/home/johannes/Documents/jina/docarrayv2/venv/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py", line 394, in _compile
    raise InternalTorchDynamoError() from e
torch._dynamo.exc.InternalTorchDynamoError

Process finished with exit code 1

Seems fixable at first glance, I'll work on it

Navigation Menu

Repository navigation

bug: torch 2.x with `TorchTensor` #1269

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

bug: torch 2.x with TorchTensor #1269

Description

Activity

h4shk4t commented on Apr 21, 2023

JohannesMessner commented on Apr 21, 2023

JohannesMessner commented on May 8, 2023

h4shk4t commented on May 10, 2023

JohannesMessner commented on May 15, 2023

JohannesMessner commented on Jul 31, 2023

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

bug: torch 2.x with `TorchTensor` #1269