<jupyter_start><jupyter_text>Multivariate derivativesOther notebooks have explained how to compute partial derivatives, but what if we want full gradients?<jupyter_code>import jax
import jax.numpy as jnp

from probfindiff import central, differentiate, from_grid, stencil<jupyter_output><empty_output><jupyter_text>Let's define a function $f: R^d \rightarrow R$.<jupyter_code>f = lambda z: jnp.dot(z, z)
d = 4

x = jnp.arange(1.0, 1.0 + d)
df = jax.jacfwd(f)<jupyter_output>WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)<jupyter_text>We have to extend the 1d scheme to a multivariate scheme.A multivariate scheme comes with a new set of coefficients and a new grid, that is adapted to the shape of the function.<jupyter_code>scheme, xs_1d = central(dx=0.01)
xs = stencil.multivariate(xs_1d=xs_1d, shape_input=(d,))
print(xs.shape)<jupyter_output>(4, 4, 3)<jupyter_text>The shape of `xs` is deliberate.The "final" dimension of the grid must correspond to the finite-difference-weight-multiplication.The shapes in the very front must correspond to the input and output shape of the domain of the function,because we aim to match the shapes of Jax' automatic differentiation behaviour.Therefore, the "evaluation domain" (which is the axis that will be "contracted" to ``shape=()`` by the scalar-valued function) must be ``axis=-2`` (the only axis left).Let us evaluate the gradient numerically now.<jupyter_code># Firstly, batch over the FD coefficients.
# Secondly, over the input shape(s).
f_batched = jax.vmap(jax.vmap(f, in_axes=-1), in_axes=0)
fx = f_batched(x[None, :, None] + xs)
dfx, _ = differentiate(fx, scheme=scheme)
print(dfx, df(x))<jupyter_output>[2.0007117 4.0015507 6.002268  8.003107 ] [2. 4. 6. 8.]<jupyter_text>The same can be done for _any_ one-dimensional scheme.<jupyter_code>xs_1d = jnp.array([-0.1, -0.01, 0.0, 0.01, 0.1])
scheme = from_grid(xs=xs_1d)
xs = stencil.multivariate(xs_1d=xs_1d, shape_input=(d,))

fx = f_batched(x[None, :, None] + xs)
dfx, _ = differentiate(fx, scheme=scheme)
print(xs.shape, dfx, df(x))<jupyter_output>(4, 4, 5) [1.9995394 3.9993134 5.9990873 7.9988513] [2. 4. 6. 8.]
{'fraction_non_alphanumeric': 0.11977971546581001, 'fraction_numerical': 0.04956402019274897, 'mean_word_length': 5.586102719033232, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '475254', 'license': 'MIT', 'n_tokens_mistral': 881, 'n_tokens_neox': 782, 'n_words': 245, 'path': 'docs/source/notebooks/multivariate_derivatives.ipynb', 'repo_name': 'pitmonticone/probfindiff', 'size': '2179'}