Tests running on CPU

[mmanzoorTT]

Testing infrastructure is running both instances on the CPU. The desired behaviour is to run one instance on tt device and other on cpu to compare the outcome of both instances. It seems that DeviceConnecter class is not behaving as expected and is unable to register tt plugin/device.

[mmanzoorTT]

@mrakitaTT @tapspatel

[tapspatel]

Repro test case

doesn't work

pytest -ssv /code/tt-xla/tests/TTIR/test_basic_ops.py::test_abs_op

works

import pytest
import jax
import jax._src.xla_bridge as xb
import jax.numpy as jnp
import numpy
import os

path = os.path.join(os.path.dirname(__file__), "../../../build/src/tt/pjrt_plugin_tt.so")
if not os.path.exists(path):
    raise FileNotFoundError(
        f"Could not find tt_pjrt C API plugin at {path}, have you compiled the project?"
    )

plugin = xb.register_plugin("tt", priority=500, library_path=path, options=None)
jax.config.update("jax_platforms", "tt,cpu")

from infrastructure import verify_module

def test_encoder_linear_one():
    def module(a, b, c, d):
        act = jnp.dot(a, b)
        act = jnp.multiply(act, c)
        return jnp.add(act, d)

    verify_module(module, input_shapes=[(1, 784), (784, 128), (1, 128), (1, 128)])

Essentially, if you force the "initialize" stage that happens in tt-xla/tests/conftest.py before importing anything else, the test will register tt-device and run on it.

[mrakitaTT]

@mmanzoorTT @tapspatel Thank you for reporting this, IIRC @kmitrovicTT was aware of this issue and should be fixed now that transition to the new test infra is finished. @kmitrovicTT can you please verify/confirm?

[mrakitaTT]

@ajakovljevicTT also mentioned this issue here: #145

[kmitrovicTT]

Sorry, missed this in my emails.

Yup, this should be gone now that there is no mix between the old and new infra. Upon running any test, plugin will be automatically registered due to the way singleton DeviceConnector is devised and instantiated.

I'm closing this issue as obsolete, should have closed it with #136.