JAX: Jit Autograd XLA

unlike , doesn’t promote for efficiency and compatibility (TPUs), it does also support

use to disable implicit type promotion (still allows python scalars type promotion for convenience)

wrap input args → record all operations

ns → produce a jaxpr using the tracer records

traces a function during the first call, wrap objects & functions in and keep track of shapes, types (like graph guards in torch compile) to produce operations that can be passed to to compute

can be used to extract the produced sequence of operations

static operations are evaluated at compile-time in python

traced operations are evaluated at run-tim in XLA

arguments to be used in control flow should be flagged as static using

or

decorator cache a compiled version of the function using it’s hash code, that as long as the traced inputs have the same type, shape will be used, else func will be recompiled and cached

when a jitted function is decorated with the cached version of the inner function will be used

avoid producing a new hash each call

arguments passed to a jitted function can be “dnoated” i.e we use the same device memory allocated for the input values of that argument to store the output / updated value (post func)

using we can trace inputs to figure out shapes (useful to avoid compiling twice when one of the arguments will start as None then will be an array)

trace and transform a given function and return a function that given a set of inputs, differentiate w.r.t to the first one

by default a transformed function compute the gradient of it’s first argument ( or ) by tracing all operations inside that function

is similar to grad, but evaluate the function & it’s grad

is similar to but operate on vector valued functions

forward mode auto-diff to compute the Jacobian vector product

i.e: directional derivative along an input vector v

v is a tangent vector in input space

forward mode vector Jacobian product: computes $\mathbf{J}.\mathbf{v}$: directional derivative along the input vector $\mathbf{v}$​

c is a cotangent vector in output space

compute the full Jacobian matrix in forward mode diff: efficient when input dim is small

compute the full Jacobian matrix in forward mode diff: efficient when output dim is small

a high level wrapper that dynamically chooses or depending on input, output shapes

evaluate a function at point / input vector v and return the output and a linear approximate of

used for efficient compute

computes the hessian matrix

alternatives:

traces inputs arguments just like and adds a batch dimension at axis 0 by default, batch dimension position can be configured for inputs and outputs with &

(if set to None, the corresponding argument won’t be vectorized)

note: & in … can be with the same structure on the corresponding input & output

and a whole branch of a can be set to the same value (no need to specify the whole structure + the same value at each leaf)

vectorized map

parallel map: SPMD

: Jax exPRession (torch IR eq)

is a representation of a sequence of primitive operations encoding the shape & type of each operation, args, inputs, outputs

a tree is a collection of tensors such as a model’s params, an optimizer state, …

a leaf’s key path is a list of keys of length = depth of leaf

provide a reader-friendly path key with a preset of keys for predefined nodes and using in custom ones

skims a collection and extract it’s structure / leaves

lax is similar to jax.numpy, by a raw, lower-level XLA primitives focus mirror of it, requires careful use, enables more efficient & general operations

at core all lax operations are python wrapper of XLA operations, with few to no guards, checks or constraints (optimize at your own risk ye 3abit)

python’s print will only work during tracing,

to keep printing / debugging in subsequent runs use

opens an interactive window at the specified point

it calls a python function from inside a transformed function

in the global random state is updated after each generation

which wouldn’t be compatible with what’s aiming to achieve while being reproducible mainly for being stateful, order depending ( execute vectorized and parallelized operations in no deterministic order)

in the random generation state / key is explicitly handled

is used to generate per-device random keys starting from the same root key injected with the device idx

until we overcome unemployment

across shards op + replicate

moves an array from host to device according to a given sharding

kwarg is used to indicate where / how to shard the result of an operation lowered in xla

xla computations are delegated to device(s) asynchronously, only shape, dtype, sharding are retrieved instantly so that the Python code can “run ahead of time”

in a training loop (fetch batch, compute) both calls overlap achieving a pseudo parallel run of both, introducing a blocking operation will ruin the orchestration

logging metrics for example (fetch batch, compute, log) would force block the fetching op until the compute operation is completed introducing a bottleneck, to mitigate that we can log the previous metrics

printing computations outside of the main, jit-compiled “compute” call would force host-device synchronization

⇒ perform additional computations (learning rate update, metrics / logs computes) using or on host in to avoid the host-device roundtrip / sync

overlaps data transfer and training on accelerators given a python iterator yielding a dict at each step (: async host → device data transfer)

a like function that could be called niside a ted function to shard intermediate / output arrays manually

is a wrapper that track an axis name, across which the array is sharded

partial replication across device computations are replicated

sharding is explicitly handled

a function is map sharded (as in CUDA kernels but for blocks) for across shards operations use [lax utils.collective operations]

a mesh is the set / structure of devices, naming axis and setting their types (Explicit, Auto, …) than the partitions an array across devices

example:

example:

input is sharded 4@x

output is sharded 4@x

each device does it’s computation that the output is sharded accordingly

a parallelization transform that execture a given function on sharded inputs as specified by then collects / re-assemble outputs as specified by on the given device mesh

(note for my fish memory: does the sharding according to not need (generally) to pre-shard inputs unless intended)

None → replicate, axis_name → shard on this mesh axis

None → fetch a single copy, no need to tile it’s already replicated (unvarying), axis_name → collect

use to label a varying value to execute operations with other varying values (mostly added in s to pass to the XLA, mentioned just in case a khazoo9 is encountered)