Index

A

accretion() (in module graphax.examples.easy) ada_lif() (in module graphax.examples.neuromorphic) ADALIF_SNN() (in module graphax.examples.neuromorphic) add_elemental_rule() (in module graphax.primitives) append_pre_transforms() (in module graphax.core)

apply() (graphax.primitives.JacobianTransform method) apply_inverse() (graphax.primitives.JacobianTransform method) atan2_elemental_rule() (in module graphax.primitives) attn() (in module graphax.examples.deep_learning) autoconversion() (in module graphax.examples.easy)

B

Black76() (in module graphax.examples.economics) BlackScholes() (in module graphax.examples.economics)

BlackScholes_Jacobian() (in module graphax.examples.economics) broadcast_elemental_rule() (in module graphax.primitives)

C

C() (in module graphax.examples.differential_kinematics) CloudSchemes_step() (in module graphax.examples.easy) concatenate_elemental_rule() (in module graphax.primitives) condensation() (in module graphax.examples.easy)

convert_element_type_rule() (in module graphax.primitives) copy() (graphax.sparse.tensor.SparseTensor method) count_muls() (in module graphax.sparse.utils) count_muls_jaxpr() (in module graphax.sparse.utils)

D

decoder_block() (in module graphax.examples.deep_learning) defelemental() (in module graphax.primitives) defelemental2() (in module graphax.primitives) dense() (graphax.sparse.tensor.SparseTensor method)

DenseDimension (class in graphax.sparse.tensor) device_put_elemental_rule() (in module graphax.primitives) dims (graphax.sparse.tensor.SparseTensor attribute) div_elemental_rule() (in module graphax.primitives) dot_general_elemental_rule() (in module graphax.primitives)

E

Encoder() (in module graphax.examples.deep_learning) encoder_block() (in module graphax.examples.deep_learning) EncoderDecoder() (in module graphax.examples.deep_learning) enthalpy_1d() (in module graphax.examples.roe) enthalpy_3d() (in module graphax.examples.roe)

eq_elemental_rule() (in module graphax.primitives) eqx_jacve() (in module graphax.equinox_bindings) evaporation() (in module graphax.examples.easy) eye_like() (in module graphax.sparse.utils) eye_like_copy() (in module graphax.sparse.utils)

F

f() (in module graphax.examples.randoms) F_1d() (in module graphax.examples.roe) F_3d() (in module graphax.examples.roe) fi() (in module graphax.examples.minpack)

filter_jacve() (in module graphax.equinox_bindings) fL() (in module graphax.examples.minpack) fR() (in module graphax.examples.minpack) FreeEnergy() (in module graphax.examples.easy)

G

g() (in module graphax.examples.randoms) GaussianDataFitting() (in module graphax.examples.minpack) get_aval_shape() (in module graphax.primitives) get_largest_tensor() (in module graphax.sparse.utils) get_ndim() (in module graphax.primitives) get_num_adds() (in module graphax.sparse.tensor) get_num_muls() (in module graphax.sparse.tensor) get_shape() (in module graphax.primitives) graphax module graphax.core module graphax.equinox_bindings module graphax.examples module graphax.examples.deep_learning module graphax.examples.differential_kinematics module

graphax.examples.easy module graphax.examples.economics module graphax.examples.minpack module graphax.examples.neuromorphic module graphax.examples.randoms module graphax.examples.roe module graphax.primitives module graphax.sparse module graphax.sparse.tensor module graphax.sparse.utils module

H

Helmholtz() (in module graphax.examples.easy)

Hole() (in module graphax.examples.easy) HumanHeartDipole() (in module graphax.examples.minpack)

I

id (graphax.sparse.tensor.DenseDimension attribute) (graphax.sparse.tensor.SparseDimension attribute)

inverse_transform (graphax.primitives.JacobianTransform attribute) iota_elemental_rule() (in module graphax.primitives)

J

JacobianTransform (class in graphax.primitives)

jacve() (in module graphax.core)

K

KerrSenn_Jacobian() (in module graphax.examples.easy)

KerrSenn_metric() (in module graphax.examples.easy)

L

layer_norm() (in module graphax.examples.deep_learning) lif() (in module graphax.examples.neuromorphic)

LIF_SNN() (in module graphax.examples.neuromorphic) Lighthouse() (in module graphax.examples.easy)

M

make_parallel_jacobian() (in module graphax.primitives) max_elemental_rule() (in module graphax.primitives) min_elemental_rule() (in module graphax.primitives) module graphax graphax.core graphax.equinox_bindings graphax.examples graphax.examples.deep_learning graphax.examples.differential_kinematics graphax.examples.easy graphax.examples.economics graphax.examples.minpack graphax.examples.neuromorphic graphax.examples.randoms graphax.examples.roe graphax.primitives graphax.sparse graphax.sparse.tensor graphax.sparse.utils

muL() (in module graphax.examples.minpack) mul_elemental_rule() (in module graphax.primitives) muR() (in module graphax.examples.minpack)

O

other_id (graphax.sparse.tensor.SparseDimension attribute)

out_dims (graphax.sparse.tensor.SparseTensor attribute)

P

Perceptron() (in module graphax.examples.deep_learning) Phi() (in module graphax.examples.economics) pjit_elemental_rule() (in module graphax.primitives) post_transforms (graphax.sparse.tensor.SparseTensor attribute) pow_elemental_rule() (in module graphax.primitives)

pre_transforms (graphax.sparse.tensor.SparseTensor attribute) prepend_post_transforms() (in module graphax.core) pressure_1d() (in module graphax.examples.roe) pressure_3d() (in module graphax.examples.roe) primal_dims (graphax.sparse.tensor.SparseTensor attribute) PropaneCombustion() (in module graphax.examples.minpack)

R

reduce_elemental_rule() (in module graphax.primitives) reduce_max_elemental_rule() (in module graphax.primitives) reduce_min_elemental_rule() (in module graphax.primitives) reduce_sum_elemental_rule() (in module graphax.primitives)

reshape_elemental_rule() (in module graphax.primitives) RobotArm_6DOF() (in module graphax.examples.differential_kinematics) RoeFlux_1d() (in module graphax.examples.roe) RoeFlux_3d() (in module graphax.examples.roe)

S

S() (in module graphax.examples.differential_kinematics) select_elemental_rule() (in module graphax.primitives) shape (graphax.sparse.tensor.SparseTensor attribute) sigmoid() (in module graphax.examples.deep_learning) SiLU() (in module graphax.examples.deep_learning) Simple() (in module graphax.examples.easy) size (graphax.sparse.tensor.DenseDimension attribute) (graphax.sparse.tensor.SparseDimension attribute) slice_elemental_rule() (in module graphax.primitives) softmax_cross_entropy() (in module graphax.examples.deep_learning)

sparse_tensor_zeros_like() (in module graphax.sparse.tensor) SparseDimension (class in graphax.sparse.tensor) SparseTensor (class in graphax.sparse.tensor) squeeze_elemental_rule() (in module graphax.primitives) standard_elemental() (in module graphax.primitives) standard_elemental2() (in module graphax.primitives) SteadyStateCombustion() (in module graphax.examples.minpack) stop_gradient_elemental_rule() (in module graphax.primitives) sub_elemental_rule() (in module graphax.primitives) superspike_surrogate() (in module graphax.examples.neuromorphic)

T

transform (graphax.primitives.JacobianTransform attribute)

transpose_elemental_rule() (in module graphax.primitives) tree_allclose() (in module graphax.core)

U

unload_post_transforms() (in module graphax.core)

unload_pre_transforms() (in module graphax.core)

V

val (graphax.sparse.tensor.SparseTensor attribute) val_dim (graphax.sparse.tensor.DenseDimension attribute) (graphax.sparse.tensor.SparseDimension attribute)

variance() (in module graphax.examples.deep_learning) vertex_elimination_jaxpr() (in module graphax.core)

W

with_type_promotion() (in module graphax.primitives)

Z

zeros_like() (in module graphax.sparse.utils)