AttributeError: Can't get attribute 'assert_no_tape_batching' on <module 'pennylane.gradients.gradient_transform'

The issue you're encountering is related to compatibility between PennyLane, PyTorch, and Amazon Braket when running hybrid quantum-classical algorithms. Let's address the problems and suggest some solutions:

1. Device compatibility: When switching from 'default.qubit' to 'braket.local.qubit', you encountered issues with tensors that have grad_fn. This is likely due to differences in how gradients are handled between the default PennyLane device and the Braket device.

2. Hybrid job error: The error you're seeing in the hybrid job ('AttributeError: Can't get attribute 'assert_no_tape_batching' on <module 'pennylane.gradients.gradient_transform'') suggests that there might be a version mismatch or compatibility issue between the PennyLane version installed in your local environment and the one available in the hybrid job environment.

To address these issues, try the following:

1. Use the 'device' differentiation method: Instead of using 'diff_method="best"', try specifying 'diff_method="device"' in your QNode definition. This tells PennyLane to use the device's native differentiation method, which may be more compatible with the Braket device.

2. Ensure consistent versions: Make sure you're using compatible versions of PennyLane, PyTorch, and the Amazon Braket SDK in both your local environment and the hybrid job environment. You may need to specify these versions explicitly in your hybrid job dependencies.

3. Use the adjoint differentiation method: When using the Braket local simulator or SV1, you can use the adjoint differentiation method. Modify your QNode definition as follows:

@qml.qnode(self.dev, interface='torch', diff_method='device')
def qnode(inputs, weights):
    # ... (rest of your code)

4. Check Braket device configuration: Ensure that your Braket device is properly configured, especially when running on the local simulator. You might need to specify the device ARN or other parameters.

5. Simplify for debugging: If the issues persist, try simplifying your quantum circuit and gradually add complexity back in. This can help isolate where the problem is occurring.

6. Use explicit typing: In your hybrid job function, consider explicitly typing your input parameters. This can sometimes help with compatibility issues.

7. Update your hybrid job decorator: Ensure you're using the latest syntax for the hybrid job decorator and specifying all necessary dependencies.

If these steps don't resolve the issue, you may need to consult the latest documentation for PennyLane, PyTorch, and Amazon Braket, as there could be recent changes or known issues that are affecting your specific use case. Additionally, reaching out to the AWS support or the PennyLane community forums might provide more specific guidance for your particular setup.
Sources
Use PennyLane with Amazon Braket - Amazon Braket

(Repeating this comment to your question since you don't seem to have noticed it)

Hello,

Could you share some more complete code that would enable use to test what you have, e.g. a minimal example that also includes actually executing the hybrid job?

Additionally, can you share the environment you are using? In particular, which version of PennyLane are you using? The Braket job containers are currently on 0.38, see https://github.com/amazon-braket/amazon-braket-containers/blob/080a1b3a54619655834c5261d7db19929f3e81d3/pytorch/jobs/docker/2.2/py3/requirements.txt#L19. It looks like assert_no_tape_batching has not existed since 0.36, so I wonder if this is a version mismatch.

Finally, is there any reason that PennyLane simulators are insufficient for your purposes (e.g. you plan to continue to hardware)?

Thank you very much for reaching me out.

I have actually fixed the issues by decorating inside @hybrid_job decorator the entire code:

@hybrid_job(device="local:braket/default", dependencies=["torch"])
def run_job():
    import numpy as np
    import time
    import pennylane as qml
    import torch
    import torch.nn as nn
    import torch.optim as optim
    from torch.utils.data import DataLoader, TensorDataset
    from braket.jobs.metrics import log_metric
    from braket.jobs import hybrid_job
    
    # Definition of the QML NN
    class QmlNN(nn.Module):
        def __init__(self, input_size, n_layers, n_qubits):

            super(QmlNN, self).__init__()

            # *Define of the quantum circuit*
            self.dev = qml.device('braket.local.qubit', wires=n_qubits)  # todo: change types of qubit and gradient

            @qml.qnode(self.dev, interface='torch', diff_method='best')
            def qnode(inputs, weights):
                inputs = inputs.detach().clone()
                # Circuit definition
                # *Amplitude data Embedding (added random padding)
                qml.templates.AmplitudeEmbedding(inputs, wires=range(n_qubits), normalize=True)  # *Pad with constant float   
                # todo: read strongly entangling layers and embedding (https://arxiv.org/abs/1804.00633)                 
                qml.templates.StronglyEntanglingLayers(weights, wires=range(n_qubits))
                return [qml.expval(qml.PauliZ(i)) for i in range(n_qubits)]

            # self.dev = qml.device('braket.local.qubit', wires=n_qubits)  # todo: change types of qubit and gradient
            # self.q_circuit = qml.QNode(qnode, self.dev, interface='torch', diff_method='best')

            # Network architecture
            self.layerIn = nn.Linear(input_size, input_size, dtype=torch.float64)   # Input layer
            self.leaky_relu = nn.LeakyReLU(negative_slope=0.3)  # LeakyRelu
            q_weights_shape = {'weights': (n_layers, n_qubits, 3)}
            # self.qlayer = qml.qnn.TorchLayer(self.q_circuit, q_weights_shape)   # *Quanutm Layer (See qnode)* 
            self.qlayer = qml.qnn.TorchLayer(qnode, q_weights_shape)
            self.layerOut = nn.Linear(n_qubits, 1, dtype=torch.float64)  # Output Layer

        # Forward pass
        def forward(self, x):
            x = self.layerIn(x)             # Data in input
            x = self.leaky_relu(x)          # Leaky relu input
            x = self.qlayer(x)              # Forward qlayer
            x = self.layerOut(x)            # Data in output
            return x
    
    # Define the training loop
    def training_model(model, criterion, optimizer, train_loader, test_loader, num_epochs):

        losses = []
        test_losses = []
        epoch_times = []

        for epoch in range(num_epochs):
            epoch_start_time = time.time()

            # Model in training mode
            model.train()
            running_loss = 0.0

            ib = 0
            for batch_x, batch_y in train_loader:
                print(f'Batch {ib}')
                optimizer.zero_grad()
                outputs = model(batch_x)
                mse_loss = criterion(outputs, batch_y.view(-1, 1))
                rmse_loss = torch.sqrt(mse_loss)
                rmse_loss.backward()
                # print("After backward:", model.qlayer.weights.grad)  # Should now have values
                optimizer.step()
                running_loss += rmse_loss.item()
                ib += 1

            # Synchronize GPU and CPU and stop timing
            epoch_end_time = time.time()
            epoch_duration = epoch_end_time - epoch_start_time
            epoch_times.append(epoch_duration)

            # Model in validation mode
            model.eval()
            test_loss = 0.0
            with torch.no_grad():
                for test_x, test_y in test_loader:
                    test_outputs = model(test_x)
                    mse_test_loss = criterion(test_outputs, test_y.view(-1, 1))
                    rmse_test_loss = torch.sqrt(mse_test_loss)
                    test_loss += rmse_test_loss.item()

        print(f'Epoch [{epoch+1}/{num_epochs}], Training Loss: {running_loss/len(train_loader):.4f}, \
                 Test Loss: {test_loss/len(test_loader):.4f}, Time: {epoch_duration:.2f} seconds')

            # Storing loss and val_loss
        losses.append(running_loss/len(train_loader))
        test_losses.append(test_loss/len(test_loader))
    
        log_metric(metric_name="test_rmse", iteration_number=epoch, value=test_loss)

        # save_model_loss(model, losses, test_losses, epoch_times)
        print('Finished Training')
        return losses, test_losses, epoch_times
    
    # Hyperparameters quantum NN
    np.random.seed(42)
    input_size = X_train.shape[1]
    num_epochs = 100

    n_layers = 1
    n_qubits = 8 # For AmplitudeEmbedding

    # Hyperparameters optimizer
    batch_size = 2
    learning_rate = 0.001

    # Train, validation and test data loader
    train_dataset = TensorDataset(X_train, y_train)
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_dataset = TensorDataset(X_test, y_test)
    test_loader = DataLoader(test_dataset, batch_size=len(test_dataset), shuffle=False)

    # Instantiate the classical model
    model_qml = QmlNN(input_size, n_layers, n_qubits)

    # Loss and optimizer
    criterion = nn.MSELoss()  # Loss function, you can choose based on your problem
    optimizer = optim.Adam(model_qml.parameters(), lr=learning_rate)
    
    losses, test_losses, times = training_model(model_qml, criterion, optimizer, train_loader, test_loader, num_epochs)
    return losses, test_losses, times

Now I am trying to run the same code on QPU (IQM Garnet and Rigetti Ankaa-3). But I got other errors: pennylane.QuantumFunctionError: Device <LegacyDeviceFacade: <Bra ketAwsQubitDevice device (wires=8, shots=500) at 0x7fe9551359c0>> does not support device with requested circuit. Is there an easy way to compile the circuit to make it work on quantum devices?

Hi Paolo,

According to the PennyLane documentation for theAttributeEncoding template, the features argument is in general not differentiable, and you should not expect to be able to optimize over it. I would suggest trying to rewrite your use case so you don't have to optimize over it.

Additionally, the error you provided seems to suggest that PennyLane has trouble choosing a good differentiation method for hardware. I suggest that you specifically set diff_method to be a hardware compatible option, e.g. "parameter-shift". See the documentation on the PennyLane website.

Dear Cedric,

The code works smoothly with the default.qubit pennylane device. I have also managed to make it work on the braket.local.qubit simulator (though it is much slower). And also on some QPUs: it is currently running on arn:aws:braket:us-west-1::device/qpu/rigetti/Ankaa-3 and arn:aws:braket:us-east-1::device/qpu/ionq/Forte-1. To train the QNN, I am printing the batch fed into the network for the gradient calculation and training. However, for some reason, the code stuck in the training of the first batch. I am wondering if this is due to some queuing priority policy, or to the performance of the QPUs, or some bugs in my code. Is it possible to print some additional details of the device? Like the chosen gradient scheme (parameter shift, adjoint, etc., in the device specification I have just indicated best), the details of the transpilation (the wiring of the QPU) etc.?

Thank you very much for help and assistance

Dear Paolo,

First, although you didn't ask for this, I'd like to mention that since you are training QNNs on QPUs, I recommend looking into the expected costs first. After a back-of-the-envelope estimate of the cost, to be on the safe side I also recommend setting a hard maximum on your number of iterations per training while also setting up checkpointing for your hybrid job, so you can pick up training later if you wish.

To gain more visibility into the progress of the training on a QPU, it's best to monitor the status of the tasks associated with each particular hybrid job, as well as the queue position of the job and the tasks. The easiest way to do this is via the Braket console or the Braket Python SDK.

To your last question: The parameters that Pennylane uses to create the Braket tasks are not logged on the Braket service. You need to log them on your side if this information is necessary for post-processing and debugging. If you'd like to attach custom labels to your jobs or tasks, tagging can be useful, but there is no simple way to instruct PennyLane to make use of this feature at the moment.

Hope this helps.