jax_functions

Override machinery and jax.numpy mirroring wrappers.

jnp.reshape(custom_obj, ...) cannot be intercepted by a custom array-like object: JAX has no override protocol, __jax_array__ is no longer honored, and pytree registration only covers jax.jit, jax.vmap and jax.tree. This module supplies the missing piece: a __jax_function__ hook with has_jax_function() and handle_jax_function(), plus thin wrappers over the jax.numpy functions probly needs. Call jax_reshape(x, ...) instead of jnp.reshape(x, ...) so custom objects get a say.

Dispatch duck-types on __jax_function__ being defined on the argument’s type, so this module never imports probly.representation.jax_like and the import graph stays one-way.

The jax_function_dispatch() decorator selects operands by parameter name and preserves public signatures and documentation. Operator overloads give native operands precise array result types; arbitrary SupportsJaxFunction overrides have an object result type. Their native implementation bodies run only after dispatch has excluded custom overrides, which is why narrowing casts to native operand types are safe there.

Classes

SupportsJaxFunction

Structural interface for custom JAX function overrides.

Functions

handle_jax_function	Dispatch a call to the __jax_function__ implementation of the arguments.
has_jax_function	Check whether any of the arguments overrides the jax function protocol.
jax_absolute	Compute absolute values elementwise, mirroring jax.numpy.absolute.
jax_add	Add two operands elementwise, mirroring jax.numpy.add.
jax_astype	Cast the array to a new data type, mirroring jax.numpy.astype.
jax_average	Compute a possibly weighted average, mirroring jax.numpy.average.
jax_bitwise_and	Compute bitwise AND, mirroring jax.numpy.bitwise_and.
jax_bitwise_or	Compute bitwise OR, mirroring jax.numpy.bitwise_or.
jax_bitwise_xor	Compute bitwise exclusive OR, mirroring jax.numpy.bitwise_xor.
jax_concatenate	Join a sequence of arrays along an existing axis, mirroring jax.numpy.concatenate.
jax_conj	Return the complex conjugate of the array, mirroring jax.numpy.conj.
jax_copy	Return a copy of the array, mirroring jax.numpy.copy.
jax_divmod	Compute quotients and remainders, mirroring jax.numpy.divmod.
jax_equal	Compare operands for equality, mirroring jax.numpy.equal.
jax_expand_dims	Insert axes of length one, mirroring jax.numpy.expand_dims.
jax_floor_divide	Floor-divide operands, mirroring jax.numpy.floor_divide.
jax_function_dispatch	Decorate a native implementation with JAX function override dispatch.
jax_greater	Test whether the left operand is larger, mirroring jax.numpy.greater.
jax_greater_equal	Test the greater-than-or-equal relation, mirroring jax.numpy.greater_equal.
jax_invert	Invert bits elementwise, mirroring jax.numpy.invert.
jax_left_shift	Shift bits left elementwise, mirroring jax.numpy.left_shift.
jax_less	Test whether the left operand is smaller, mirroring jax.numpy.less.
jax_less_equal	Test the less-than-or-equal relation, mirroring jax.numpy.less_equal.
jax_matmul	Multiply matrices or vectors, mirroring jax.numpy.matmul.
jax_matrix_transpose	Transpose the last two axes of the array, mirroring jax.numpy.matrix_transpose.
jax_mean	Compute the arithmetic mean, mirroring jax.numpy.mean.
jax_moveaxis	Move axes of the array to new positions, mirroring jax.numpy.moveaxis.
jax_multiply	Multiply operands elementwise, mirroring jax.numpy.multiply.
jax_negative	Negate an operand elementwise, mirroring jax.numpy.negative.
jax_not_equal	Compare operands for inequality, mirroring jax.numpy.not_equal.
jax_positive	Apply unary positive, mirroring jax.numpy.positive.
jax_power	Raise operands to powers elementwise, mirroring jax.numpy.power.
jax_remainder	Compute floor-division remainders, mirroring jax.numpy.remainder.
jax_reshape	Return a reshaped version of the array, mirroring jax.numpy.reshape.
jax_right_shift	Shift bits right elementwise, mirroring jax.numpy.right_shift.
jax_squeeze	Remove axes of length one, mirroring jax.numpy.squeeze.
jax_stack	Join a sequence of arrays along a new axis, mirroring jax.numpy.stack.
jax_std	Compute the standard deviation, mirroring jax.numpy.std.
jax_subtract	Subtract operands elementwise, mirroring jax.numpy.subtract.
jax_sum	Compute the sum, mirroring jax.numpy.sum.
jax_swapaxes	Swap two axes of the array, mirroring jax.numpy.swapaxes.
jax_take_along_axis	Take values along an axis, mirroring jax.numpy.take_along_axis.
jax_transpose	Return a transposed version of the array, mirroring jax.numpy.transpose.
jax_true_divide	Divide operands elementwise, mirroring jax.numpy.true_divide.
jax_var	Compute the variance, mirroring jax.numpy.var.
try_jax_function	Try function overrides without preventing Python's reflected operator fallback.