torch.GradientsFor

export type GradientsFor<Op extends OpName> = OpSchemas[Op] extends BinaryOpSchema
  ? { a: TensorLike | null; b: TensorLike | null }
  : OpSchemas[Op] extends UnaryOpSchema
    ? { input: TensorLike | null }
    : OpSchemas[Op] extends ReductionOpSchema
      ? { input: TensorLike | null }
      : { input: TensorLike | null };

Opextends OpName

Derive gradient return type from operation schema.

Binary operations return gradients for both inputs (a, b). Unary and reduction operations return gradient for single input.