Not working with current stable covalent

Covalent version: 0.220.0.post2 Python version: 3.8

Following error is raised when we try to run with the latest braket image:

Traceback (most recent call last):
  File "/opt/ml/code/exec.py", line 24, in <module>
    result = function(*args, **kwargs)
  File "/usr/local/lib/python3.8/site-packages/covalent/executor/base.py", line 106, in wrapper_fn
    output = fn(*new_args, **new_kwargs)
  File "/home/avalanche/anaconda3/envs/qa-38/lib/python3.8/site-packages/covalent_dispatcher/_core/runner.py", line 260, in qelectron_compatible_wrapper
  File "/tmp/ipykernel_14208/1216327594.py", line 29, in simple_quantum_task
  File "/usr/local/lib/python3.8/site-packages/pennylane/__init__.py", line 329, in device
    plugin_device_class = plugin_devices[name].load()
  File "/usr/local/lib/python3.8/site-packages/pkg_resources/__init__.py", line 2450, in load
    return self.resolve()
  File "/usr/local/lib/python3.8/site-packages/pkg_resources/__init__.py", line 2456, in resolve
    module = __import__(self.module_name, fromlist=['__name__'], level=0)
  File "/usr/local/lib/python3.8/site-packages/braket/pennylane_plugin/__init__.py", line 14, in <module>
    from braket.pennylane_plugin.ahs_device import (  # noqa: F401
  File "/usr/local/lib/python3.8/site-packages/braket/pennylane_plugin/ahs_device.py", line 50, in <module>
    from .ahs_translation import (
  File "/usr/local/lib/python3.8/site-packages/braket/pennylane_plugin/ahs_translation.py", line 133, in <module>
    def _create_register(coordinates: list[tuple[float, float]]):
TypeError: 'type' object is not subscriptable

The workflow was used to test is:

import covalent as ct
import os

# AWS resources to pass to the executor
braket_job_execution_role_name = "<obtained from tf deployment>"
ecr_image_uri = "<obtained from tf deployment>"
s3_bucket_name = "<obtained from tf deployment>"

# Instantiate the executor
ex = ct.executor.BraketExecutor(
            s3_bucket_name=s3_bucket_name,
            ecr_image_uri=ecr_image_uri,
            braket_job_execution_role_name=braket_job_execution_role_name,
    )

# Execute the following circuit:
# |0> - H - Measure
@ct.electron(executor=ex)
def simple_quantum_task(num_qubits: int):
    import pennylane as qml

    # These are passed to the Hybrid Jobs container at runtime
    device_arn = os.environ["AMZN_BRAKET_DEVICE_ARN"]
    s3_bucket = os.environ["AMZN_BRAKET_OUT_S3_BUCKET"]
    s3_task_dir = os.environ["AMZN_BRAKET_TASK_RESULTS_S3_URI"].split(s3_bucket)[1]

    device = qml.device(
        "braket.aws.qubit",
        device_arn=device_arn,
        s3_destination_folder=(s3_bucket, s3_task_dir),
        wires=num_qubits,
    )

    @qml.qnode(device=device)
    def simple_circuit():
        qml.Hadamard(wires=[0])
        return qml.expval(qml.PauliZ(wires=[0]))

    res = simple_circuit().numpy()
    return res

@ct.lattice
def simple_quantum_workflow(num_qubits: int):
    return simple_quantum_task(num_qubits=num_qubits)

dispatch_id = ct.dispatch(simple_quantum_workflow)(1)
print(dispatch_id)

AgnostiqHQ / covalent-braket-plugin

Not working with current stable covalent #76