class GraphSageLoss extends SingleParentVariable<Scalar> {

    @Override
    public Matrix gradient(Variable<?> parent, ComputationContext ctx) {
        Tensor<?> embeddingData = ctx.data(parent);
        double[] embeddings = embeddingData.data();
        int totalBatchSize = embeddingData.dimension(ROWS_INDEX);
        int batchSize = totalBatchSize / 3;

        int embeddingSize = embeddingData.dimension(COLUMNS_INDEX);
        double[] gradientResult = new double[totalBatchSize * embeddingSize];

        IntStream.range(0, batchSize).forEach(nodeId -> {
            int positiveNodeId = nodeId + batchSize;
            int negativeNodeId = nodeId + NEGATIVE_NODES_OFFSET * batchSize;
            int dimension = parent.dimension(COLUMNS_INDEX);
            double positiveAffinity = affinity(embeddingData, nodeId, positiveNodeId);
            double negativeAffinity = affinity(embeddingData, nodeId, negativeNodeId);

            double positiveLogistic = logisticFunction(positiveAffinity);
            double negativeLogistic = logisticFunction(-negativeAffinity);

            IntStream.range(0, embeddingSize).forEach(columnOffset -> partialComputeGradient(
                embeddings,
                gradientResult,
                nodeId,
                positiveNodeId,
                negativeNodeId,
                dimension,
                positiveLogistic,
                negativeLogistic,
                columnOffset
            ));

        });
        return new Matrix(gradientResult, totalBatchSize, embeddingSize);
    }

      GraphSageLoss(Variable<Matrix> combinedEmbeddings, int negativeSamplingFactor);

}