Skip to main content

A jax wrapper for autograd-differentiable functions.

Project description

Agjax -- jax wrapper for autograd-differentiable functions.

v0.3.3

Agjax allows existing code built with autograd to be used with the jax framework.

In particular, agjax.wrap_for_jax allows arbitrary autograd functions ot be differentiated using jax.grad. Several other function transformations (e.g. compilation via jax.jit) are not supported.

Meanwhile, agjax.experimental.wrap_for_jax supports grad, jit, vmap, and jacrev. However, it depends on certain under-the-hood behavior by jax, which is not guaranteed to remain unchanged. It also is more restrictive in terms of the valid function signatures of functions to be wrapped: all arguments and outputs must be convertible to valid jax types. (agjax.wrap_for_jax also supports non-jax inputs and outputs, e.g. strings.)

Installation

pip install agjax

Usage

Basic usage is as follows:

@agjax.wrap_for_jax
def fn(x, y):
  return x * npa.cos(y)

jax.grad(fn, argnums=(0,  1))(1.0, 0.0)

# (Array(1., dtype=float32), Array(0., dtype=float32))

The experimental wrapper is similar, but requires that the function outputs and datatypes be specified, simiilar to jax.pure_callback.

wrapped_fn = agjax.experimental.wrap_for_jax(
  lambda x, y: x * npa.cos(y),
  result_shape_dtypes=jnp.ones((5,)),
)

jax.jacrev(wrapped_fn, argnums=0)(jnp.arange(5, dtype=float), jnp.arange(5, 10, dtype=float))

# [[ 0.28366217  0.          0.          0.          0.        ]
#  [ 0.          0.96017027  0.          0.          0.        ]
#  [ 0.          0.          0.75390226  0.          0.        ]
#  [ 0.          0.          0.         -0.14550003  0.        ]
#  [ 0.          0.          0.          0.         -0.91113025]]

Agjax wrappers are intended to be quite general, and can support functions with multiple inputs and outputs as well as functions that have nondifferentiable outputs or arguments that cannot be differentiated with respect to. These should be specified using nondiff_argnums and nondiff_outputnums arguments. In the experimental wrapper, these must still be jax-convertible types, while in the standard wrapper they may have arbitrary types.

@functools.partial(
  agjax.wrap_for_jax, nondiff_argnums=(2,), nondiff_outputnums=(1,)
)
def fn(x, y, string_arg):
  return x * npa.cos(y), string_arg * 2

(value, aux), grad = jax.value_and_grad(
  fn, argnums=(0, 1), has_aux=True
)(1.0, 0.0, "test")

print(f"value = {value}")
print(f"  aux = {aux}")
print(f" grad = {grad}")
value = 1.0
  aux = testtest
 grad = (Array(1., dtype=float32), Array(0., dtype=float32))

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

agjax-0.3.3.tar.gz (11.2 kB view hashes)

Uploaded Source

Built Distribution

agjax-0.3.3-py3-none-any.whl (9.7 kB view hashes)

Uploaded Python 3

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page