Diffractor
Diffractor.AbstractTangentBundleDiffractor.ExplicitTangentDiffractor.JetDiffractor.TangentBundleDiffractor.TaylorTangentDiffractor.UniformTangentDiffractor.zero_bundleDiffractor.∂xⁿDiffractor.∂⃖Diffractor.∂☆Diffractor.∇Diffractor.:'Diffractor.bundleDiffractor.destructureDiffractor.dxDiffractor.find_end_of_phi_blockDiffractor.find_taylor_incompatibilityDiffractor.forward_diff_no_inf!Diffractor.jet_taylor_evDiffractor.@∂
Diffractor.AbstractTangentBundle — Typeabstract type AbstractTangentBundle{N, B}; endThis type represents the N-th order (iterated) tangent bundle [1] TⁿB over some base (Riemannian) manifold B. Note that TⁿB is itself another manifold and thus in particular a vector space (over ℝ). As such, subtypes of this abstract type are expected to support the usual vector space operations.
However, beyond that, this abstract type makes no guarantee about the representation. That said, to gain intution for what this object is, it makes sense to pick some explicit bases and write down examples.
To that end, suppose that B=ℝ. Then T¹B=T¹ℝ is just our usual notion of a dual number, i.e. for some element η ∈ T¹ℝ, we may consider η = a + bϵ for real numbers (a, b) and ϵ an infinitessimal differential such that ϵ^2 = 0.
Equivalently, we may think of η as being identified with the vector (a, b) ∈ ℝ² with some additional structure. The same form essentially holds for general B, where we may write (as sets):
T¹B = {(a, b) | a ∈ B, b ∈ Tₐ B }Note that these vectors are orthogonal to those in the underlying base space. For example, if B=ℝ², then we have:
T¹ℝ² = {([aₓ, a_y], [bₓ, b_y]) | [aₓ, a_y] ∈ ℝ², [bₓ, b_y] ∈ Tₐ ℝ² }For convenience, we will sometimes writes these in one as:
η ∈ T ℝ² = aₓ x̂ + a_y ŷ + bₓ ∂/∂x|_aₓ x + b_y ∂/∂y|_{a_y}
+ := aₓ x̂ + a_y ŷ + bₓ ∂/∂¹x + b_y ∂/∂¹y
+ := [aₓ, a_y] + [bₓ, b_y] ∂/∂¹
+ := a + b ∂/∂¹
+ := a + b ∂₁These are all definitional equivalences and we will mostly work with the final form. An important thing to keep in mind though is that the subscript on ∂₁ does not refer to a dimension of the underlying base manifold (for which we will rarely pick an explicit basis here), but rather tags the basis of the tangent bundle.
Let us iterate this construction to second order. We have:
T²B = T¹(T¹B) = { (α, β) | α ∈ T¹B, β ∈ T_α T¹B }
+ = { ((a, b), (c, d)) | a ∈ B, b ∈ Tₐ B, c ∈ Tₐ B, d ∈ T²ₐ B}(where in the last equality we used the linearity of the tangent vector).
Following our above notation, we will canonically write such an element as:
a + b ∂₁ + c ∂₂ + d ∂₂ ∂₁
+= a + b ∂₁ + c ∂₂ + d ∂₁ ∂₂It is worth noting that there still only is one base point a of the underlying manifold and thus TⁿB is a vector bundle over B for all N.
Further Reading
[1] https://en.wikipedia.org/wiki/Tangent_bundle
Diffractor.ExplicitTangent — Typestruct ExplicitTangent{P}A fully explicit coordinate representation of the tangent space, represented by a vector of 2^N-1 partials.
Diffractor.Jet — Typestruct Jet{T, N}Represents the truncated (N-1)-th order Taylor series
f(a) + (x-a)f'(a) + 1/2(x-a)^2f''(a) + ...Coefficients are stored in unscaled form. For a jet j, several operations are supported:
- Indexing
j[i]returns fᵢ - Evaluation j(x) semantically evaluates the truncated taylor series at
x. However, evaluation is restricted to be precisely ata- the additional information in the taylor series is only available through derivatives. Mathematically this corresponds to an infinitessimal ball arounda.
Diffractor.TangentBundle — Typestruct TangentBundle{N, B, P}Represents a tangent bundle as an explicit primal together with some representation of the tangent space.
Diffractor.TaylorTangent — Typestruct TaylorTangent{C}The taylor bundle construction mods out the full N-th order tangent bundle by the equivalence relation that coefficients of like-order basis elements be equal, i.e. rather than a generic element
a + b ∂₁ + c ∂₂ + d ∂₃ + e ∂₂ ∂₁ + f ∂₃ ∂₁ + g ∂₃ ∂₂ + h ∂₃ ∂₂ ∂₁we have a tuple (c₀, c₁, c₂, c₃) corresponding to the full element
c₀ + c₁ ∂₁ + c₁ ∂₂ + c₁ ∂₃ + c₂ ∂₂ ∂₁ + c₂ ∂₃ ∂₁ + c₂ ∂₃ ∂₂ + c₃ ∂₃ ∂₂ ∂₁i.e.
c₀ + c₁ (∂₁ + ∂₂ + ∂₃) + c₂ (∂₂ ∂₁ + ∂₃ ∂₁ + ∂₃ ∂₂) + c₃ ∂₃ ∂₂ ∂₁This restriction forms a submanifold of the original manifold. The naming is by analogy with the (truncated) Taylor series
c₀ + c₁ x + 1/2 c₂ x² + 1/3! c₃ x³ + O(x⁴)Diffractor.UniformTangent — Typestruct UniformTangentRepresents an N-th order tangent bundle with all unform partials. Particularly useful for representing singleton values.
Diffractor.zero_bundle — Type(::zero_bundle{N})(primal)Creates a bundle with a zero tangent.
Diffractor.∂xⁿ — Type∂xⁿ{N}(x)For x in a one dimensional manifold, map x to the trivial, unital, Nth order tangent bundle. It should hold that ∀x ⟨∂ⁿ{1}x(x), dx(x)⟩ = 1
Diffractor.∂⃖ — Type∂⃖{N}∂⃖{N} is the reverse-mode AD optic functor of order N. A call (::∂⃖{N})(f, args...) corresponds to ∂⃖ⁿ f(args...) in the linear encoding of an N-optic (see the terminology guide for definitions of these terms).
In general (::∂⃖{N})(f, args...) will return a tuple of the original primal value f(args...) (in rare cases primitives may modify the primal value - in general we will ignore this rare complication for the purposes of clear documentation) and an optic continuation λ. The interpretation of this continuation depends on the order of the functor:
For example, ∂⃖{1} computes first derivatives. In particular, for a function f, ∂⃖{1}(f, args...) will return the tuple (f(args...), f⋆) (read "f upper-star").
Diffractor.∂☆ — Type∂☆{N,E}∂☆{N} is the forward-mode AD functor of order N (An integer). A call (::∂☆{N})(f, args...) evaluating a function f: A -> B is lifted to its pushforward on the N-th order tangent bundle f⋆: Tⁿ A -> Tⁿ B.
!!!advanced "Eras Mode" E (a bool, default false) is for Eras mode. In Eras mode, we are Taylor or bust. Normally if a particular derivative can not be represented as a TaylorBundle we fall back and represent it as a ExplictTangentBundle. However, in Eras mode we error if it can't be represented as a TaylorBundle. In general, this is not wanted since it often will break nested AD. But in the cases it doesn't its really fast, since it means we can rewrite nested AD as Taylor-mode AD (plus its more type stable). To be safe in Eras mode, it is sufficient, but not necessary, to be doing nested AD with respect to the same variable. It also works in other cases where (likely by problem construction) ADing with respect to a second variable happens to result in something that can be represented with a TaylorBundle also. (You need your different partials to happen to be exactly equal).
Diffractor.∇ — Type∇(f, args...)Computes the gradient ∇f(x, y, z...) (at (x, y, z...)). In particular, the return value will be a tuple of partial derivatives (∂f/∂x, ∂f/∂y, ∂f/∂z...).
Curried version
Alternatively, ∇ may be curried, essentially giving the gradient as a function:
Examples
julia> using Diffractor: ∇
+
+julia> map(∇(*), (1,2,3), (4,5,6))
+((4.0, 1.0), (5.0, 2.0), (6.0, 3.0))The derivative ∂f/∂f
Note that since in Julia, there is no distinction between functions and values, there is in principle a partial derivative with respect to the function itself. However, said partial derivative is dropped by this interface. It is however available using the lower level ∂⃖ if desired. This interaction can also be used to obtain gradients with respect to only some of the arguments by using a closure:
∇((x,z)->f(x,y,z))(x, z) # returns (∂f/∂x, ∂f/∂z)
Though of course the same can be obtained by simply indexing the resulting tuple (in well-inferred code there should not be a performance difference between these two options).
Diffractor.:' — Methodf'This is a convenience syntax for taking the derivative of a function f: ℝ -> ℝ. In particular, for such a function f'(x) will be the first derivative of f at x (and similar for f''(x) and second derivatives and so on.)
Note that the syntax conflicts with the Base definition for the adjoint of a matrix and thus is not enabled by default. To use it, add the following to the top of your module:
using Diffractor: var"'"It is also available using the @∂ macro:
@∂ f'(x)Diffractor.bundle — Functionbundle(primal, tangent)Wraps a primal up with a tangent into the appropriate kind of AbstractBundle{1}. This is more or less the Diffractor equivelent of ForwardDiff.jl's Dual type.
Diffractor.destructure — Methodfor a TaylorTangent{N, <:Tuple} this breaks it up unto 1 TaylorTangent{N} for each element of the primal tuple
Diffractor.dx — Methoddx(x)dx represents the trival differential one-form of a one dimensional Riemannian manifold M. In particular, it is a section of the cotangent bundle of M, meaning it may be evaluted at a point x of M to obtain an element of the cotangent space T*ₓ M to M at x. We impose no restrictions on the representations of either the manifold itself or the cotangent space.
By default, the only implementation provided identifies T*ₓ ℝ ≃ ℝ, keeping watever type is used to represent ℝ. i.e.
dx(x::Real) = one(x)However, users may provide additional overloads for custom representations of one dimensional Riemannian manifolds.
Diffractor.find_end_of_phi_block — Methodfind_end_of_phi_block(ir::IRCode, start_search_idx::Int)Finds the last index within the same basic block, on or after the start_search_idx which is not within a phi block. A phi-block is a run on PhiNodes or nothings that must be the first statements within the basic block.
If start_search_idx is not within a phi block to begin with, then just returns start_search_idx
Diffractor.find_taylor_incompatibility — Methodfinds the lowerest order derivative that is not taylor compatible, or returns -1 if all compatible
Diffractor.forward_diff_no_inf! — Methodforward_diff_no_inf!(ir::IRCode, to_diff::Vector{Pair{SSAValue,Int}}; visit_custom!, transform!)Internal method which generates the code for forward mode diffentiation
irthe IR being differnetationto_diff: collection of all SSA values for which the derivative is to be taken, paired with the order (first deriviative, second derivative etc)visit_custom!(ir::IRCode, ssa, order::Int, recurse::Bool) -> Bool:
decides if the custom `transform!` should be applied to a `stmt` or not
+ Default: `false` for all statementstransform!(ir::IRCode, ssa::SSAValue, order::Int)mutatesirto do a custom tranformation.eras_mode: determines if to error if not all derivatives are taylor
Diffractor.jet_taylor_ev — Methodjet_taylor_ev(::Val{}, jet, taylor)Generates a closed form arithmetic expression for the N-th component of the action of a 1d jet (of order at least N) on a maximally symmetric (i.e. taylor) tangent bundle element. In particular, if we represent both the jet and the taylor tangent bundle element by their associated canonical taylor series:
j = j₀ + j₁ (x - a) + j₂ 1/2 (x - a)^2 + ... + jₙ 1/n! (x - a)^n
+t = t₀ + t₁ (x - t₀) + t₂ 1/2 (x - t₀)^2 + ... + tₙ 1/n! (x - t₀)^nthen the action of evaluating j on t, is some other taylor series
t′ = a + t′₁ (x - a) + t′₂ 1/2 (x - a)^2 + ... + t′ₙ 1/n! (x - a)^nThe t′ᵢ can be found by explicitly plugging in t for every x and expanding out, dropping terms of orders that are higher. This computes closed form expressions for the t′ᵢ that are hopefully easier on the compiler.
Diffractor.@∂ — Macro@∂Convenice macro for writing partial derivatives. E.g. The expression:
@∂ f(∂x, ∂y)Will compute the partial derivative ∂^2 f/∂x∂y at (x, y)`. And similarly
@∂ f(∂²x, ∂y)will compute the derivative ∂^3 f/∂x^2 ∂y at (x,y).