code: def train_return_grad(self, input, grad):
          
  
          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
docs: 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.