Pointwise operators

PyTorch's core ATen vocabulary contains many pointwise operators, such as neg, div, pow, rsqrt, erf, gelu and clamp_min. They share one shape rule: every operand and the result have the same shape. Graphs therefore represent them as one operation family whose operator is data, and PointwiseOp lists the operators of that family together with the static attributes PyTorch passes them, such as the rounding mode of torch.div.

What an operator computes depends on the scalar type. A scalar dictionary lists the operators it implements in Context.pointwise, as a PointwiseTable. Each entry is a total function that returns the result or a reason why the scalar type has none; over the reals, for example, a zero divisor has no quotient. An operator missing from a table is unsupported for that scalar type, and evaluators report it instead of approximating it.

References

• PyTorch Core ATen IR: https://docs.pytorch.org/docs/stable/torch.compiler_ir.html
• torch.div: https://docs.pytorch.org/docs/stable/generated/torch.div.html
• torch.nn.functional.gelu: https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.gelu.html

inductive Spec.DivRounding : Type

Rounding applied to a quotient, as the rounding_mode argument of torch.div.

• none : DivRounding

  True division, PyTorch's rounding_mode=None.

• trunc : DivRounding

  Round the quotient toward zero, rounding_mode="trunc".

• floor : DivRounding

  Round the quotient toward negative infinity, rounding_mode="floor".

def Spec.instReprDivRounding.repr : DivRounding → Nat → Std.Format

@[instance_reducible]

instance Spec.instReprDivRounding : Repr DivRounding

@[instance_reducible]

instance Spec.instDecidableEqDivRounding : DecidableEq DivRounding

@[instance_reducible]

instance Spec.instInhabitedDivRounding : Inhabited DivRounding

inductive Spec.GeluApproximation : Type

The formula selected by the approximate argument of torch.nn.functional.gelu.

• none : GeluApproximation

  x Φ(x), with the standard normal distribution function Φ written through erf.

• tanh : GeluApproximation

  The tanh approximation 0.5 x (1 + tanh (√(2 / π) (x + 0.044715 x³))).

@[instance_reducible]

instance Spec.instReprGeluApproximation : Repr GeluApproximation

def Spec.instReprGeluApproximation.repr : GeluApproximation → Nat → Std.Format

@[instance_reducible]

instance Spec.instDecidableEqGeluApproximation : DecidableEq GeluApproximation

@[instance_reducible]

instance Spec.instInhabitedGeluApproximation : Inhabited GeluApproximation

inductive Spec.PointwiseOp : Type

Pointwise operators represented by the graph family NN.IR.OpKind.pointwise.

• neg : PointwiseOp

  Negation, torch.neg.

• rsqrt : PointwiseOp

  Reciprocal square root, torch.rsqrt.

• erf : PointwiseOp

  Gauss error function, torch.erf.

• gelu (approximate : GeluApproximation) : PointwiseOp

  Gaussian error linear unit, torch.nn.functional.gelu.

• div (rounding : DivRounding) : PointwiseOp

  Division with an optional rounding of the quotient, torch.div.

• pow : PointwiseOp

  Power with a tensor exponent, torch.pow.

• clampMin : PointwiseOp

  Lower bound with a scalar bound, torch.clamp_min: the bound where the operand is below it, and a NaN where either is one.

@[instance_reducible]

instance Spec.instReprPointwiseOp : Repr PointwiseOp

def Spec.instReprPointwiseOp.repr : PointwiseOp → Nat → Std.Format

@[instance_reducible]

instance Spec.instDecidableEqPointwiseOp : DecidableEq PointwiseOp

def Spec.instDecidableEqPointwiseOp.decEq (x✝ x✝¹ : PointwiseOp) : Decidable (x✝ = x✝¹)

@[instance_reducible]

instance Spec.instInhabitedPointwiseOp : Inhabited PointwiseOp

def Spec.PointwiseOp.arity : PointwiseOp → Nat

Number of tensor operands.

theorem Spec.PointwiseOp.arity_pos (op : PointwiseOp) : 0 < op.arity

Every operator takes at least one operand.

def Spec.PointwiseOp.name : PointwiseOp → String

Short operator name without attributes, for grouping diagnostics.

def Spec.PointwiseOp.describe : PointwiseOp → String

Operator name with its attribute, for diagnostics. Artifact codecs define their own spelling.

def Spec.PointwiseOp.all : List PointwiseOp

Every operator with every attribute value, in declaration order.

theorem Spec.PointwiseOp.mem_all (op : PointwiseOp) : op ∈ all

The enumeration is complete.

@[reducible, inline]

abbrev Spec.PointwiseFn (α : Type) (arity : Nat) : Type

The meaning of an operator with arity operands at the scalar type α.

The function receives the operands in order and returns the result, or the reason why the scalar type has no result for them.

def Spec.PointwiseFn.unary {α : Type} (f : α → Except String α) : PointwiseFn α 1

A meaning with one operand.

def Spec.PointwiseFn.binary {α : Type} (f : α → α → Except String α) : PointwiseFn α 2

A meaning with two operands, in order.

@[reducible, inline]

abbrev Spec.PointwiseTable (α : Type) : Type

The pointwise operators a scalar type implements, each with its meaning.

def Spec.divFloorSchedule {α : Type} [Zero α] [One α] [Add α] [Sub α] [Div α] [LT α] [DecidableRel fun (x1 x2 : α) => x1 < x2] [BEq α] (fmod : α → α → α) (floor signedZero : α → α) (half a b : α) : α

PyTorch's flooring division, c10::div_floor_floating, over the operations it uses.

fmod is C's truncating remainder, floor rounds toward negative infinity, and signedZero q is zero with the sign of q. A zero divisor returns the quotient itself. Otherwise the remainder corrects the quotient, as Python's // does, and the final comparison with half repairs a quotient that rounded across an integer. A binary format runs this schedule with its rounded operations; over the reals it computes ⌊a / b⌋.