TFE_Op* MatMulOp(TFE_Context* ctx, TFE_TensorHandle* a, TFE_TensorHandle* b);

// Return an identity op.
TFE_Op* IdentityOp(TFE_Context* ctx, TFE_TensorHandle* a);

// Return a shape op fetching the shape of `a`.
TFE_Op* ShapeOp(TFE_Context* ctx, TFE_TensorHandle* a);

// Return an allreduce op adding up input tensor `in` from `group_size` workers.

    for (auto accumulated : accumulated_gradients_) {
      vspace_.DeleteGradient(accumulated.second);
    }
  }

  // Tell the forward accumulator to watch tensor_id, with a Tensor tangent
  // vector `tangent` of matching shape and dtype. Tangents are the "vector" in
  // "Jacobian-vector product"; `Watch`ing a new Tensor and immediately calling
  // FetchJVP for it would return `tangent`.
  void Watch(int64_t tensor_id, Gradient* tangent);