Error when using torch.block_diag method

[joelamzn]

The following code throws the below error on trn1.32xlarge instance.

>>> import torch, torch_xla
>>> import torch_xla.core.xla_model as xm
>>> device = xm.xla_device()
>>> segments = [torch.ones((1,4), device=device)]
>>> torch.block_diag(*segments[:])

Error:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/local/lib/python3.8/dist-packages/torch/functional.py", line 1266, in block_diag
    return torch._C._VariableFunctions.block_diag(tensors)  # type: ignore[attr-defined]
RuntimeError: torch_xla/csrc/aten_xla_type.cpp:3426 : Check failed: !runtime::sys_util::GetEnvBool("XLA_DISABLE_FUNCTIONALIZATION", false)
*** Begin stack trace ***
    tsl::CurrentStackTrace()
    torch_xla::XLANativeFunctions::block_diag(c10::ArrayRef<at::Tensor>)

    at::_ops::block_diag::redispatch(c10::DispatchKeySet, c10::ArrayRef<at::Tensor>)

    at::_ops::block_diag::call(c10::ArrayRef<at::Tensor>)

    PyCFunction_Call
    _PyObject_MakeTpCall
    _PyEval_EvalFrameDefault
    _PyEval_EvalCodeWithName
    _PyFunction_Vectorcall
    PyObject_Call
    _PyEval_EvalFrameDefault
    _PyEval_EvalCodeWithName
    PyEval_EvalCode

    PyRun_InteractiveLoopFlags
    PyRun_AnyFileExFlags

    Py_BytesMain
    __libc_start_main
    _start
*** End stack trace ***

Setup

• Torch: 2.1.2+cu121
• Torch XLA: 2.1.1

[jeffhataws]

Hi @joelamzn ,

Thanks for reporting the issue. Currently we disable functionalization as default for performance reason. Will you try setting XLA_DISABLE_FUNCTIONALIZATION=0 to run your example. I tried your code with this environment variable setting and no longer see the error.

[aws-taylor]

Hello @joelamzn,

We haven't heard from you in a while, so I'm going to resolve this issue. Feel free to re-open if you require further asstance.