// (where the value starts on the host), without allowing other implicit
        // copies/broadcasts. Other implicit copies may be supported eventually,
        // but need special handling for gradients (gradient of copy-on is not
        // just copy-off but includes a sum) and consideration of performance.
        //
        // TODO(allenl): There may be smarter ways to do this copy in some

        }
      }
    } else {
      (*result)[id].push_back(output_gradients[i]);
    }
  }
  return absl::OkStatus();
}

// TODO(agarwal): use an automatic mechanism for handling None arguments to
// gradient functions.
//
// Some gradient functions can accept None arguments for gradients. The
// following maps the operation name to the indices at which the corresponding