Performs a forward pass from the input batch, followed by a backward
        pass using the provided gradient (in place of a loss function). Returns
        a MpsFloatArray representing the output (final gradient) of the backward
        pass. Calling asnumpy() on this value will wait for the batch to finish
        and yield the output as a numpy array.

def train_return_grad(self, input, grad):
        """
        Performs a forward pass from the input batch, followed by a backward
        pass using the provided gradient (in place of a loss function). Returns
        a MpsFloatArray representing the output (final gradient) of the backward
        pass. Calling asnumpy() on this value will wait for the batch to finish
        and yield the output as a numpy array.
        """

        assert self._mode == MpsGraphMode.TrainReturnGrad
        assert input.shape == self._ishape
        assert grad.shape == self._oshape

        input_array = MpsFloatArray(input)
        grad_array = MpsFloatArray(grad)
        result_handle = _ctypes.c_void_p()
        status_code = self._LIB.TCMPSTrainGraph(
            self.handle, input_array.handle, grad_array.handle,
                _ctypes.byref(result_handle))

        assert status_code == 0, "Error calling TCMPSTrainReturnGradGraph"
        assert result_handle, "TCMPSTrainReturnGradGraph unexpectedly returned NULL pointer"

        result = MpsFloatArray(result_handle)
        assert result.shape() == self._ishape

        return result