Documentation

LeanMlir.Proofs.Codegen.StableHLO.PrecisionSwitch

Precision-switched constructors — `XAt bf16 rnd …` is `XBf16 rnd …` or `X …` #

Every bf16 op is its f32 peer with one extra leading argument, the rounding rnd. XAt bf16 rnd … is the if bf16 then .XBf16 rnd … else .X … choice, written once per constructor instead of at every call site. pretty evaluates it, so the emitted text is exactly the chosen branch's.

A leaf on Basic: the renderers (RenderKit and the *Render* modules) and the typed forward graphs in Nets/ (r34IdGraphB, …) both choose the constructor here, so a render and the graph its faithfulness theorem is about cannot pick different kinds for the same flag. The renderers pass zrnd = id (Pretty.lean), the graphs pass id, and Foundation.Bf16Erasure states each switch at id equal to its f32 peer (denOp_convAt_id, den_convBackBatchedAt_id, …).

@[reducible]

def Proofs.StableHLO.BatchableOp.convAt (bf16 : Bool) {ic oc h w kH kW : ℕ} (rnd : ℝ → ℝ) (wName bName : String) (W : Kernel4 oc ic kH kW) (bias : Vec oc) :

BatchableOp (ic * h * w) (oc * h * w)

conv, or its bf16 peer at rounding rnd when bf16.

Equations

One or more equations did not get rendered due to their size.

Instances For

@[reducible]

def Proofs.StableHLO.SHlo.convBackBatchedAt (bf16 : Bool) {N ic oc h w kH kW : ℕ} (rnd : ℝ → ℝ) (wName : String) (W : Kernel4 oc ic kH kW) (b : Vec oc) :

SHlo (N * (oc * h * w)) → SHlo (N * (ic * h * w))

convBackBatched, or its bf16 peer at rounding rnd when bf16.

Equations

Proofs.StableHLO.SHlo.convBackBatchedAt bf16 rnd wName W b = if bf16 = true then Proofs.StableHLO.SHlo.convBackBatchedBf16 rnd wName W b else Proofs.StableHLO.SHlo.convBackBatched wName W b

Instances For

@[reducible]

def Proofs.StableHLO.BatchableOp.convStride4At (bf16 : Bool) {ic oc h w kH kW : ℕ} (rnd : ℝ → ℝ) (wName bName : String) (W : Kernel4 oc ic kH kW) (bias : Vec oc) :

BatchableOp (ic * (2 * (2 * h)) * (2 * (2 * w))) (oc * h * w)

convStride4, or its bf16 peer at rounding rnd when bf16.

Equations

One or more equations did not get rendered due to their size.

Instances For

@[reducible]

def Proofs.StableHLO.SHlo.convStride4WeightGradBAt (bf16 : Bool) {N ic oc h w kH kW : ℕ} (rnd : ℝ → ℝ) (xName : String) (b : Vec oc) (x : Vec (N * (ic * (2 * (2 * h)) * (2 * (2 * w))))) (W : Kernel4 oc ic kH kW) :

SHlo (N * (oc * h * w)) → SHlo (oc * ic * kH * kW)

convStride4WeightGradB, or its bf16 peer at rounding rnd when bf16.

Equations

One or more equations did not get rendered due to their size.

Instances For

@[reducible]

def Proofs.StableHLO.BatchableOp.convStridedAt (bf16 : Bool) {ic oc h w kH kW : ℕ} (rnd : ℝ → ℝ) (wName bName : String) (W : Kernel4 oc ic kH kW) (bias : Vec oc) :

BatchableOp (ic * (2 * h) * (2 * w)) (oc * h * w)

convStrided, or its bf16 peer at rounding rnd when bf16.

Equations

One or more equations did not get rendered due to their size.

Instances For

@[reducible]

def Proofs.StableHLO.SHlo.convStridedBackBatchedAt (bf16 : Bool) {N ic oc h w kH kW : ℕ} (rnd : ℝ → ℝ) (wName : String) (W : Kernel4 oc ic kH kW) (b : Vec oc) :

SHlo (N * (oc * h * w)) → SHlo (N * (ic * (2 * h) * (2 * w)))

convStridedBackBatched, or its bf16 peer at rounding rnd when bf16.

Equations

One or more equations did not get rendered due to their size.

Instances For

@[reducible]

def Proofs.StableHLO.SHlo.convStridedWeightGradBAt (bf16 : Bool) {N ic oc h w kH kW : ℕ} (rnd : ℝ → ℝ) (xName : String) (b : Vec oc) (x : Vec (N * (ic * (2 * h) * (2 * w)))) (W : Kernel4 oc ic kH kW) :

SHlo (N * (oc * h * w)) → SHlo (oc * ic * kH * kW)

convStridedWeightGradB, or its bf16 peer at rounding rnd when bf16.

Equations

One or more equations did not get rendered due to their size.

Instances For

@[reducible]

def Proofs.StableHLO.BatchableOp.convStridedXlaAt (bf16 : Bool) {ic oc h w kH kW : ℕ} (rnd : ℝ → ℝ) (wName bName : String) (W : Kernel4 oc ic kH kW) (bias : Vec oc) :

BatchableOp (ic * (2 * h) * (2 * w)) (oc * h * w)

convStridedXla, or its bf16 peer at rounding rnd when bf16.

Equations

One or more equations did not get rendered due to their size.

Instances For

@[reducible]

def Proofs.StableHLO.SHlo.convStridedXlaWeightGradBAt (bf16 : Bool) {N ic oc h w kH kW : ℕ} (rnd : ℝ → ℝ) (xName : String) (b : Vec oc) (x : Vec (N * (ic * (2 * h) * (2 * w)))) (W : Kernel4 oc ic kH kW) :

SHlo (N * (oc * h * w)) → SHlo (oc * ic * kH * kW)

convStridedXlaWeightGradB, or its bf16 peer at rounding rnd when bf16.

Equations

One or more equations did not get rendered due to their size.

Instances For

@[reducible]

def Proofs.StableHLO.SHlo.convWeightGradBAt (bf16 : Bool) {N ic oc h w kH kW : ℕ} (rnd : ℝ → ℝ) (xName : String) (b : Vec oc) (x : Vec (N * (ic * h * w))) (W : Kernel4 oc ic kH kW) :

SHlo (N * (oc * h * w)) → SHlo (oc * ic * kH * kW)

convWeightGradB, or its bf16 peer at rounding rnd when bf16.

Equations

Proofs.StableHLO.SHlo.convWeightGradBAt bf16 rnd xName b x W = if bf16 = true then Proofs.StableHLO.SHlo.convWeightGradBBf16 rnd xName b x W else Proofs.StableHLO.SHlo.convWeightGradB xName b x W

Instances For

@[reducible]

def Proofs.StableHLO.BatchableOp.denseRowAt (bf16 : Bool) {N a c : ℕ} (rnd : ℝ → ℝ) (wName bName : String) (W : Mat a c) (b : Vec c) :

BatchableOp (N * a) (N * c)

denseRow, or its bf16 peer at rounding rnd when bf16.

Equations

One or more equations did not get rendered due to their size.

Instances For

@[reducible]

def Proofs.StableHLO.BatchableOp.denseRowBackAt (bf16 : Bool) {rows a c : ℕ} (rnd : ℝ → ℝ) (wName : String) (W : Mat a c) :

BatchableOp (rows * c) (rows * a)

denseRowBack, or its bf16 peer at rounding rnd when bf16.

Equations

Proofs.StableHLO.BatchableOp.denseRowBackAt bf16 rnd wName W = if bf16 = true then Proofs.StableHLO.BatchableOp.denseRowBackBf16 rnd wName W else Proofs.StableHLO.BatchableOp.denseRowBack wName W

Instances For

@[reducible]

def Proofs.StableHLO.BatchableOp.depthwiseAt (bf16 : Bool) {c h w kH kW : ℕ} (rnd : ℝ → ℝ) (wName bName : String) (W : DepthwiseKernel c kH kW) (bias : Vec c) :

BatchableOp (c * h * w) (c * h * w)

depthwise, or its bf16 peer at rounding rnd when bf16.

Equations

One or more equations did not get rendered due to their size.

Instances For

@[reducible]

def Proofs.StableHLO.SHlo.depthwiseBackBatchedAt (bf16 : Bool) {N c h w kH kW : ℕ} (rnd : ℝ → ℝ) (wName : String) (W : DepthwiseKernel c kH kW) (b : Vec c) :

SHlo (N * (c * h * w)) → SHlo (N * (c * h * w))

depthwiseBackBatched, or its bf16 peer at rounding rnd when bf16.

Equations

One or more equations did not get rendered due to their size.

Instances For

@[reducible]

def Proofs.StableHLO.BatchableOp.depthwiseStridedAt (bf16 : Bool) {c h w kH kW : ℕ} (rnd : ℝ → ℝ) (wName bName : String) (W : DepthwiseKernel c kH kW) (bias : Vec c) :

BatchableOp (c * (2 * h) * (2 * w)) (c * h * w)

depthwiseStrided, or its bf16 peer at rounding rnd when bf16.

Equations

One or more equations did not get rendered due to their size.

Instances For

@[reducible]

def Proofs.StableHLO.SHlo.depthwiseStridedBackBatchedAt (bf16 : Bool) {N c h w kH kW : ℕ} (rnd : ℝ → ℝ) (wName : String) (W : DepthwiseKernel c kH kW) (b : Vec c) :

SHlo (N * (c * h * w)) → SHlo (N * (c * (2 * h) * (2 * w)))

depthwiseStridedBackBatched, or its bf16 peer at rounding rnd when bf16.

Equations

One or more equations did not get rendered due to their size.

Instances For

@[reducible]

def Proofs.StableHLO.SHlo.depthwiseStridedWeightGradBAt (bf16 : Bool) {N c h w kH kW : ℕ} (rnd : ℝ → ℝ) (xName : String) (b : Vec c) (x : Vec (N * (c * (2 * h) * (2 * w)))) (W : DepthwiseKernel c kH kW) :

SHlo (N * (c * h * w)) → SHlo (c * kH * kW)

depthwiseStridedWeightGradB, or its bf16 peer at rounding rnd when bf16.

Equations

One or more equations did not get rendered due to their size.

Instances For

@[reducible]

def Proofs.StableHLO.BatchableOp.depthwiseStridedXlaAt (bf16 : Bool) {c h w kH kW : ℕ} (rnd : ℝ → ℝ) (wName bName : String) (W : DepthwiseKernel c kH kW) (bias : Vec c) :

BatchableOp (c * (2 * h) * (2 * w)) (c * h * w)

depthwiseStridedXla, or its bf16 peer at rounding rnd when bf16.

Equations

One or more equations did not get rendered due to their size.

Instances For

@[reducible]

def Proofs.StableHLO.SHlo.depthwiseStridedXlaBackBatchedAt (bf16 : Bool) {N c h w kH kW : ℕ} (rnd : ℝ → ℝ) (wName : String) (W : DepthwiseKernel c kH kW) (b : Vec c) :

SHlo (N * (c * h * w)) → SHlo (N * (c * (2 * h) * (2 * w)))

depthwiseStridedXlaBackBatched, or its bf16 peer at rounding rnd when bf16.

Equations

One or more equations did not get rendered due to their size.

Instances For

@[reducible]

def Proofs.StableHLO.SHlo.depthwiseStridedXlaWeightGradBAt (bf16 : Bool) {N c h w kH kW : ℕ} (rnd : ℝ → ℝ) (xName : String) (b : Vec c) (x : Vec (N * (c * (2 * h) * (2 * w)))) (W : DepthwiseKernel c kH kW) :

SHlo (N * (c * h * w)) → SHlo (c * kH * kW)

depthwiseStridedXlaWeightGradB, or its bf16 peer at rounding rnd when bf16.

Equations

One or more equations did not get rendered due to their size.

Instances For

@[reducible]

def Proofs.StableHLO.SHlo.depthwiseWeightGradBAt (bf16 : Bool) {N c h w kH kW : ℕ} (rnd : ℝ → ℝ) (xName : String) (b : Vec c) (x : Vec (N * (c * h * w))) (W : DepthwiseKernel c kH kW) :

SHlo (N * (c * h * w)) → SHlo (c * kH * kW)

depthwiseWeightGradB, or its bf16 peer at rounding rnd when bf16.

Equations

One or more equations did not get rendered due to their size.

Instances For

@[reducible]

def Proofs.StableHLO.SHlo.flatConvFAt (bf16 : Bool) {ic oc h w kH kW : ℕ} (rnd : ℝ → ℝ) (wName bName : String) (W : Kernel4 oc ic kH kW) (b : Vec oc) :

SHlo (ic * h * w) → SHlo (oc * h * w)

flatConvF, or its bf16 peer at rounding rnd when bf16.

Equations

Proofs.StableHLO.SHlo.flatConvFAt bf16 rnd wName bName W b = if bf16 = true then Proofs.StableHLO.SHlo.flatConvFBf16 rnd wName bName W b else Proofs.StableHLO.SHlo.flatConvF wName bName W b

Instances For

@[reducible]

def Proofs.StableHLO.SHlo.matmulFBAt (bf16 : Bool) {N m k n : ℕ} (rnd : ℝ → ℝ) :

SHlo (N * (m * k)) → SHlo (N * (k * n)) → SHlo (N * (m * n))

matmulFB, or its bf16 peer at rounding rnd when bf16.

Equations

Proofs.StableHLO.SHlo.matmulFBAt bf16 rnd = if bf16 = true then Proofs.StableHLO.SHlo.matmulFBBf16 rnd else Proofs.StableHLO.SHlo.matmulFB

Instances For

@[reducible]

def Proofs.StableHLO.BatchableOp.patchEmbedAt (bf16 : Bool) {ic H W P N D : ℕ} (rnd : ℝ → ℝ) (wName bName clsName posName : String) (Wc : Kernel4 D ic P P) (bc cls : Vec D) (pos : Mat (N + 1) D) :

BatchableOp (ic * H * W) ((N + 1) * D)

patchEmbed, or its bf16 peer at rounding rnd when bf16.

Equations

One or more equations did not get rendered due to their size.

Instances For

@[reducible]

def Proofs.StableHLO.SHlo.patchEmbedWeightGradBAt (bf16 : Bool) {N ic H W P tk D : ℕ} (rnd : ℝ → ℝ) (xName : String) (x : Vec (N * (ic * H * W))) :

SHlo (N * ((tk + 1) * D)) → SHlo (D * ic * P * P)

patchEmbedWeightGradB, or its bf16 peer at rounding rnd when bf16.

Equations

One or more equations did not get rendered due to their size.

Instances For

@[reducible]

def Proofs.StableHLO.SHlo.rowDenseWeightGradBAt (bf16 : Bool) {N tk a c : ℕ} (rnd : ℝ → ℝ) (xName : String) (x : Vec (N * (tk * a))) :

SHlo (N * (tk * c)) → SHlo (a * c)

rowDenseWeightGradB, or its bf16 peer at rounding rnd when bf16.

Equations

Proofs.StableHLO.SHlo.rowDenseWeightGradBAt bf16 rnd xName x = if bf16 = true then Proofs.StableHLO.SHlo.rowDenseWeightGradBBf16 rnd xName x else Proofs.StableHLO.SHlo.rowDenseWeightGradB xName x

Instances For