adv_jax_math.batch_vectorize

adv_jax_math.batch_vectorize(pyfunc, *, excluded=frozenset({}), signature=None, batch_size=None, **kwargs)[source]

Define a vectorized function with broadcasting and batching.

References

The original copyright notice is as follows Copyright 2020 The JAX Authors. Licensed under the Apache License, Version 2.0 (the “License”); https://github.com/jax-ml/jax/blob/main/jax/_src/numpy/vectorize.py.

Notes

vectorize() is a convenience wrapper for defining vectorized functions with broadcasting, in the style of NumPy’s generalized universal functions. It allows for defining functions that are automatically repeated across any leading dimensions, without the implementation of the function needing to be concerned about how to handle higher dimensional inputs.

jax.numpy.vectorize() has the same interface as numpy.vectorize, but it is syntactic sugar for an auto-batching transformation (vmap()) rather than a Python loop. This should be considerably more efficient, but the implementation must be written in terms of functions that act on JAX arrays.

Parameters:

  • pyfunc (callable) – Function to vectorize.

  • excluded (set of int or str, optional) – Positional or keyword arguments that will not be vectorized. These are passed directly to pyfunc without modification.

  • signature (str, optional) – Generalized universal function signature, such as (m,n),(n)->(m) for vectorized matrix-vector multiplication. If provided, pyfunc receives and must return arrays whose shapes match the corresponding core dimensions.

  • batch_size (int, optional) – Number of mapped elements evaluated per batch. By default, all mapped elements are evaluated together, matching jax.numpy.vectorize(). chunk_size is accepted as an alias when batch_size is None.

Returns:

callable – Batch-vectorized version of the given function.