Introduction Type Safety Tensor Expressions Tensor Indexing Einsum Einops Autograd Training a Model Profiling & Memory PyTorch Migration Best Practices Runtimes Performance PyTorch Compatibility Benchmarks DType Coverage

torch.UnaryOpConfig

export interface UnaryOpConfig<Op extends UnaryOpNames> {
  /**
   * Operation name (must match OpSchemas).
   */
  readonly op: Op;

  /**
   * CPU kernel function. Receives a single value and returns the result.
   * For float32 dtype.
   */
  readonly cpu: (x: number) => number;

  /**
   * Optional CPU kernel for int32 dtype.
   * If not provided and int32 is in dtypes, uses the float32 kernel.
   */
  readonly cpuInt32?: (x: number) => number;

  /**
   * WebGPU shader entry point name (e.g., 'exp_op', 'sin_op').
   */
  readonly shaderEntry: string;

  /**
   * Dtypes to register. Defaults to ['float32'].
   */
  readonly dtypes?: readonly UnaryDType[];

  /**
   * Backward function for autograd.
   * If not provided, the operation is not differentiable.
   */
  readonly backward?: UnaryBackwardFn;

  /**
   * Setup context function for backward.
   * Defaults to saving input if backward is provided.
   */
  readonly setupContext?: (
    ctx: GradContext,
    inputs: { input: TensorLike },
    output: TensorLike
  ) => void;

  /**
   * What to save for backward pass.
   * 'input' - saves the input tensor (default)
   * 'output' - saves the output tensor (useful for exp, sigmoid, etc.)
   * 'none' - saves nothing (for ops like neg where grad doesn't need saved tensors)
   */
  readonly saveForBackward?: SaveForBackward;

  /**
   * Tensor method name to call. Defaults to the op name.
   */
  readonly tensorMethod?: string;
}

Opextends UnaryOpNames

readonlyop(Op): – Operation name (must match OpSchemas).
readonlycpu((x: number) => number): – CPU kernel function. Receives a single value and returns the result. For float32 dtype.
readonlycpuInt32((x: number) => number)optional: – Optional CPU kernel for int32 dtype. If not provided and int32 is in dtypes, uses the float32 kernel.
readonlyshaderEntry(string): – WebGPU shader entry point name (e.g., 'exp_op', 'sin_op').
readonlydtypes(readonly UnaryDType[])optional: – Dtypes to register. Defaults to ['float32'].
readonlybackward(UnaryBackwardFn)optional: – Backward function for autograd. If not provided, the operation is not differentiable.
readonlysetupContext(( ctx: GradContext, inputs: { input: TensorLike }, output: TensorLike ) => void)optional: – Setup context function for backward. Defaults to saving input if backward is provided.
readonlysaveForBackward(SaveForBackward)optional: – What to save for backward pass. 'input' - saves the input tensor (default) 'output' - saves the output tensor (useful for exp, sigmoid, etc.) 'none' - saves nothing (for ops like neg where grad doesn't need saved tensors)
readonlytensorMethod(string)optional: – Tensor method name to call. Defaults to the op name.

Configuration for registering a unary operation.