quchip.declarative.qnp

Trace-safe numeric namespace for declarative model authors.

Use this namespace inside model implementations (value(), local_hamiltonian(), interaction()) so expressions stay JAX-traceable.

exception quchip.declarative.qnp.ComplexWarning[source]

Bases: RuntimeWarning

The warning raised when casting a complex dtype to a real dtype.

As implemented, casting a complex number to a real discards its imaginary part, but this behavior may not be what the user actually wants.

quchip.declarative.qnp.abs(x, /)

Alias of jax.numpy.absolute().

Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.absolute(x, /)

Calculate the absolute value element-wise.

JAX implementation of numpy.absolute.

This is the same function as jax.numpy.abs().

Args:

x: Input array

Returns:

An array-like object containing the absolute value of each element in x, with the same shape as x. For complex valued input, \(a + ib\), the absolute value is \(\sqrt{a^2+b^2}\).

Examples:
>>> x1 = jnp.array([5, -2, 0, 12])
>>> jnp.absolute(x1)
Array([ 5,  2,  0, 12], dtype=int32)

>>> x2 = jnp.array([[ 8, -3, 1],[ 0, 9, -6]])
>>> jnp.absolute(x2)
Array([[8, 3, 1],
       [0, 9, 6]], dtype=int32)

>>> x3 = jnp.array([8 + 15j, 3 - 4j, -5 + 0j])
>>> jnp.absolute(x3)
Array([17.,  5.,  5.], dtype=float32)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.acos(x, /)

Alias of jax.numpy.arccos()

Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.acosh(x, /)

Alias of jax.numpy.arccosh()

Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.all(a, axis=None, out=None, keepdims=False, *, where=None)

Test whether all array elements along a given axis evaluate to True.

JAX implementation of numpy.all().

Args:

a: Input array. axis: int or array, default=None. Axis along which to be tested. If None,

tests along all the axes.

keepdims: bool, default=False. If true, reduced axes are left in the result

with size 1.

where: int or array of boolean dtype, default=None. The elements to be used

in the test. Array should be broadcast compatible to the input.

out: Unused by JAX.

Returns:

An array of boolean values.

Examples:

By default, jnp.all tests for True values along all the axes.

>>> x = jnp.array([[True, True, True, False],
...                [True, False, True, False],
...                [True, True, False, False]])
>>> jnp.all(x)
Array(False, dtype=bool)

If axis=0, tests for True values along axis 0.

>>> jnp.all(x, axis=0)
Array([ True, False, False, False], dtype=bool)

If keepdims=True, ndim of the output will be same of that of the input.

>>> jnp.all(x, axis=0, keepdims=True)
Array([[ True, False, False, False]], dtype=bool)

To include specific elements in testing for True values, you can use a``where``.

>>> where=jnp.array([[1, 0, 1, 0],
...                  [0, 0, 1, 1],
...                  [1, 1, 1, 0]], dtype=bool)
>>> jnp.all(x, axis=0, keepdims=True, where=where)
Array([[ True,  True, False, False]], dtype=bool)
Parameters:
Return type:

Array

quchip.declarative.qnp.allclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False)

Check if two arrays are element-wise approximately equal within a tolerance.

JAX implementation of numpy.allclose().

Essentially this function evaluates the following condition:

\[|a - b| \le \mathtt{atol} + \mathtt{rtol} * |b|\]

jnp.inf in a will be considered equal to jnp.inf in b.

Args:

a: first input array to compare. b: second input array to compare. rtol: relative tolerance used for approximate equality. Default = 1e-05. atol: absolute tolerance used for approximate equality. Default = 1e-08. equal_nan: Boolean. If True, NaNs in a will be considered

equal to NaNs in b. Default is False.

Returns:

Boolean scalar array indicating whether the input arrays are element-wise approximately equal within the specified tolerances.

See Also:

  • jax.numpy.isclose()

  • jax.numpy.equal()

Examples:
>>> jnp.allclose(jnp.array([1e6, 2e6, 3e6]), jnp.array([1e6, 2e6, 3e7]))
Array(False, dtype=bool)
>>> jnp.allclose(jnp.array([1e6, 2e6, 3e6]),
...              jnp.array([1.00008e6, 2.00008e7, 3.00008e8]), rtol=1e3)
Array(True, dtype=bool)
>>> jnp.allclose(jnp.array([1e6, 2e6, 3e6]),
...              jnp.array([1.00001e6, 2.00002e6, 3.00009e6]), atol=1e3)
Array(True, dtype=bool)
>>> jnp.allclose(jnp.array([jnp.nan, 1, 2]),
...              jnp.array([jnp.nan, 1, 2]), equal_nan=True)
Array(True, dtype=bool)
Parameters:
Return type:

Array

quchip.declarative.qnp.amax(a, axis=None, out=None, keepdims=False, initial=None, where=None)

Alias of jax.numpy.max().

Parameters:
Return type:

Array

quchip.declarative.qnp.amin(a, axis=None, out=None, keepdims=False, initial=None, where=None)

Alias of jax.numpy.min().

Parameters:
Return type:

Array

quchip.declarative.qnp.angle(z, deg=False)

Return the angle of a complex valued number or array.

JAX implementation of numpy.angle().

Args:

z: A complex number or an array of complex numbers. deg: Boolean. If True, returns the result in degrees else returns

in radians. Default is False.

Returns:

An array of counterclockwise angle of each element of z, with the same shape as z of dtype float.

Examples:

If z is a number

>>> z1 = 2+3j
>>> jnp.angle(z1)
Array(0.98279375, dtype=float32, weak_type=True)

If z is an array

>>> z2 = jnp.array([[1+3j, 2-5j],
...                 [4-3j, 3+2j]])
>>> with jnp.printoptions(precision=2, suppress=True):
...     print(jnp.angle(z2))
[[ 1.25 -1.19]
 [-0.64  0.59]]

If deg=True.

>>> with jnp.printoptions(precision=2, suppress=True):
...     print(jnp.angle(z2, deg=True))
[[ 71.57 -68.2 ]
 [-36.87  33.69]]
Parameters:
Return type:

Array

quchip.declarative.qnp.any(a, axis=None, out=None, keepdims=False, *, where=None)

Test whether any of the array elements along a given axis evaluate to True.

JAX implementation of numpy.any().

Args:

a: Input array. axis: int or array, default=None. Axis along which to be tested. If None,

tests along all the axes.

keepdims: bool, default=False. If true, reduced axes are left in the result

with size 1.

where: int or array of boolean dtype, default=None. The elements to be used

in the test. Array should be broadcast compatible to the input.

out: Unused by JAX.

Returns:

An array of boolean values.

Examples:

By default, jnp.any tests along all the axes.

>>> x = jnp.array([[True, True, True, False],
...                [True, False, True, False],
...                [True, True, False, False]])
>>> jnp.any(x)
Array(True, dtype=bool)

If axis=0, tests along axis 0.

>>> jnp.any(x, axis=0)
Array([ True,  True,  True, False], dtype=bool)

If keepdims=True, ndim of the output will be same of that of the input.

>>> jnp.any(x, axis=0, keepdims=True)
Array([[ True,  True,  True, False]], dtype=bool)

To include specific elements in testing for True values, you can use a``where``.

>>> where=jnp.array([[1, 0, 1, 0],
...                  [0, 1, 0, 1],
...                  [1, 0, 1, 0]], dtype=bool)
>>> jnp.any(x, axis=0, keepdims=True, where=where)
Array([[ True, False,  True, False]], dtype=bool)
Parameters:
Return type:

Array

quchip.declarative.qnp.append(arr, values, axis=None)

Return a new array with values appended to the end of the original array.

JAX implementation of numpy.append().

Args:

arr: original array. values: values to be appended to the array. The values must have

the same number of dimensions as arr, and all dimensions must match except in the specified axis.

axis: axis along which to append values. If None (default), both arr

and values will be flattened before appending.

Returns:

A new array with values appended to arr.

See also:

  • jax.numpy.insert()

  • jax.numpy.delete()

Examples:
>>> a = jnp.array([1, 2, 3])
>>> b = jnp.array([4, 5, 6])
>>> jnp.append(a, b)
Array([1, 2, 3, 4, 5, 6], dtype=int32)

Appending along a specific axis:

>>> a = jnp.array([[1, 2],
...                [3, 4]])
>>> b = jnp.array([[5, 6]])
>>> jnp.append(a, b, axis=0)
Array([[1, 2],
       [3, 4],
       [5, 6]], dtype=int32)

Appending along a trailing axis:

>>> a = jnp.array([[1, 2, 3],
...                [4, 5, 6]])
>>> b = jnp.array([[7], [8]])
>>> jnp.append(a, b, axis=1)
Array([[1, 2, 3, 7],
       [4, 5, 6, 8]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.apply_along_axis(func1d, axis, arr, *args, **kwargs)

Apply a function to 1D array slices along an axis.

JAX implementation of numpy.apply_along_axis(). While NumPy implements this iteratively, JAX implements this via jax.vmap(), and so func1d must be compatible with vmap.

Args:
func1d: a callable function with signature func1d(arr, /, *args, **kwargs)

where *args and **kwargs are the additional positional and keyword arguments passed to apply_along_axis().

axis: integer axis along which to apply the function. arr: the array over which to apply the function. args, kwargs: additional positional and keyword arguments are passed through

to func1d.

Returns:

The result of func1d applied along the specified axis.

See also:

  • jax.vmap(): a more direct way to create a vectorized version of a function.

  • jax.numpy.apply_over_axes(): repeatedly apply a function over multiple axes.

  • jax.numpy.vectorize(): create a vectorized version of a function.

Examples:

A simple example in two dimensions, where the function is applied either row-wise or column-wise:

>>> x = jnp.array([[1, 2, 3],
...                [4, 5, 6]])
>>> def func1d(x):
...   return jnp.sum(x ** 2)
>>> jnp.apply_along_axis(func1d, 0, x)
Array([17, 29, 45], dtype=int32)
>>> jnp.apply_along_axis(func1d, 1, x)
Array([14, 77], dtype=int32)

For 2D inputs, this can be equivalently expressed using jax.vmap(), though note that vmap specifies the mapped axis rather than the applied axis:

>>> jax.vmap(func1d, in_axes=1)(x)  # same as applying along axis 0
Array([17, 29, 45], dtype=int32)
>>> jax.vmap(func1d, in_axes=0)(x)  # same as applying along axis 1
Array([14, 77], dtype=int32)

For 3D inputs, apply_along_axis() is equivalent to mapping over two dimensions:

>>> x_3d = jnp.arange(24).reshape(2, 3, 4)
>>> jnp.apply_along_axis(func1d, 2, x_3d)
Array([[  14,  126,  366],
       [ 734, 1230, 1854]], dtype=int32)
>>> jax.vmap(jax.vmap(func1d))(x_3d)
Array([[  14,  126,  366],
       [ 734, 1230, 1854]], dtype=int32)

The applied function may also take arbitrary positional or keyword arguments, which should be passed directly as additional arguments to apply_along_axis():

>>> def func1d(x, exponent):
...   return jnp.sum(x ** exponent)
>>> jnp.apply_along_axis(func1d, 0, x, exponent=3)
Array([ 65, 133, 243], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.apply_over_axes(func, a, axes)

Apply a function repeatedly over specified axes.

JAX implementation of numpy.apply_over_axes().

Args:
func: the function to apply, with signature func(Array, int) -> Array, and

where y = func(x, axis) must satisfy y.ndim in [x.ndim, x.ndim - 1].

a: N-dimensional array over which to apply the function. axes: the sequence of axes over which to apply the function.

Returns:

An N-dimensional array containing the result of the repeated function application.

See also:

  • jax.numpy.apply_along_axis(): apply a 1D function along a single axis.

Examples:

This function is designed to have similar semantics to typical associative jax.numpy reductions over one or more axes with keepdims=True. For example:

>>> x = jnp.array([[1, 2, 3],
...                [4, 5, 6]])

>>> jnp.apply_over_axes(jnp.sum, x, [0])
Array([[5, 7, 9]], dtype=int32)
>>> jnp.sum(x, [0], keepdims=True)
Array([[5, 7, 9]], dtype=int32)

>>> jnp.apply_over_axes(jnp.min, x, [1])
Array([[1],
       [4]], dtype=int32)
>>> jnp.min(x, [1], keepdims=True)
Array([[1],
       [4]], dtype=int32)

>>> jnp.apply_over_axes(jnp.prod, x, [0, 1])
Array([[720]], dtype=int32)
>>> jnp.prod(x, [0, 1], keepdims=True)
Array([[720]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.arange(start, stop=None, step=None, dtype=None, *, device=None)

Create an array of evenly-spaced values.

JAX implementation of numpy.arange(), implemented in terms of jax.lax.iota().

Similar to Python’s range() function, this can be called with a few different positional signatures:

  • jnp.arange(stop): generate values from 0 to stop, stepping by 1.

  • jnp.arange(start, stop): generate values from start to stop, stepping by 1.

  • jnp.arange(start, stop, step): generate values from start to stop, stepping by step.

Like with Python’s range() function, the starting value is inclusive, and the stop value is exclusive.

Args:

start: start of the interval, inclusive. stop: optional end of the interval, exclusive. If not specified, then

(start, stop) = (0, start)

step: optional step size for the interval. Default = 1. dtype: optional dtype for the returned array; if not specified it will

be determined via type promotion of start, stop, and step.

device: (optional) Device or Sharding

to which the created array will be committed.

Returns:

Array of evenly-spaced values from start to stop, separated by step.

Note:

Using arange with a floating-point step argument can lead to unexpected results due to accumulation of floating-point errors, especially with lower-precision data types like float8_* and bfloat16. To avoid precision errors, consider generating a range of integers, and scaling it to the desired range. For example, instead of this:

jnp.arange(-1, 1, 0.01, dtype='bfloat16')

it can be more accurate to generate a sequence of integers, and scale them:

(jnp.arange(-100, 100) * 0.01).astype('bfloat16')
Examples:

Single-argument version specifies only the stop value:

>>> jnp.arange(4)
Array([0, 1, 2, 3], dtype=int32)

Passing a floating-point stop value leads to a floating-point result:

>>> jnp.arange(4.0)
Array([0., 1., 2., 3.], dtype=float32)

Two-argument version specifies start and stop, with step=1:

>>> jnp.arange(1, 6)
Array([1, 2, 3, 4, 5], dtype=int32)

Three-argument version specifies start, stop, and step:

>>> jnp.arange(0, 2, 0.5)
Array([0. , 0.5, 1. , 1.5], dtype=float32)
See Also:

  • jax.numpy.linspace(): generate a fixed number of evenly-spaced values.

  • jax.lax.iota(): directly generate integer sequences in XLA.

Parameters:
Return type:

Array

quchip.declarative.qnp.arccos(x, /)

Compute element-wise inverse of trigonometric cosine of input.

JAX implementation of numpy.arccos.

Args:

x: input array or scalar.

Returns:

An array containing the inverse trigonometric cosine of each element of x in radians in the range [0, pi], promoting to inexact dtype.

Note:

  • jnp.arccos returns nan when x is real-valued and not in the closed interval [-1, 1].

  • jnp.arccos follows the branch cut convention of numpy.arccos for complex inputs.

See also:

  • jax.numpy.cos(): Computes a trigonometric cosine of each element of input.

  • jax.numpy.arcsin() and jax.numpy.asin(): Computes the inverse of trigonometric sine of each element of input.

  • jax.numpy.arctan() and jax.numpy.atan(): Computes the inverse of trigonometric tangent of each element of input.

Examples:
>>> x = jnp.array([-2, -1, -0.5, 0, 0.5, 1, 2])
>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.arccos(x)
Array([  nan, 3.142, 2.094, 1.571, 1.047, 0.   ,   nan], dtype=float32)

For complex inputs:

>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.arccos(4-1j)
Array(0.252+2.097j, dtype=complex64, weak_type=True)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.arccosh(x, /)

Calculate element-wise inverse of hyperbolic cosine of input.

JAX implementation of numpy.arccosh.

The inverse of hyperbolic cosine is defined by:

\[arccosh(x) = \ln(x + \sqrt{x^2 - 1})\]
Args:

x: input array or scalar.

Returns:

An array of same shape as x containing the inverse of hyperbolic cosine of each element of x, promoting to inexact dtype.

Note:

  • jnp.arccosh returns nan for real-values in the range [-inf, 1).

  • jnp.arccosh follows the branch cut convention of numpy.arccosh for complex inputs.

See also:

  • jax.numpy.cosh(): Computes the element-wise hyperbolic cosine of the input.

  • jax.numpy.arcsinh(): Computes the element-wise inverse of hyperbolic sine of the input.

  • jax.numpy.arctanh(): Computes the element-wise inverse of hyperbolic tangent of the input.

Examples:
>>> x = jnp.array([[1, 3, -4],
...                [-5, 2, 7]])
>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.arccosh(x)
Array([[0.   , 1.763,   nan],
       [  nan, 1.317, 2.634]], dtype=float32)

For complex-valued input:

>>> x1 = jnp.array([-jnp.inf+0j, 1+2j, -5+0j])
>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.arccosh(x1)
Array([  inf+3.142j, 1.529+1.144j, 2.292+3.142j], dtype=complex64)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.arcsin(x, /)

Compute element-wise inverse of trigonometric sine of input.

JAX implementation of numpy.arcsin.

Args:

x: input array or scalar.

Returns:

An array containing the inverse trigonometric sine of each element of x in radians in the range [-pi/2, pi/2], promoting to inexact dtype.

Note:

  • jnp.arcsin returns nan when x is real-valued and not in the closed interval [-1, 1].

  • jnp.arcsin follows the branch cut convention of numpy.arcsin for complex inputs.

See also:

  • jax.numpy.sin(): Computes a trigonometric sine of each element of input.

  • jax.numpy.arccos() and jax.numpy.acos(): Computes the inverse of trigonometric cosine of each element of input.

  • jax.numpy.arctan() and jax.numpy.atan(): Computes the inverse of trigonometric tangent of each element of input.

Examples:
>>> x = jnp.array([-2, -1, -0.5, 0, 0.5, 1, 2])
>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.arcsin(x)
Array([   nan, -1.571, -0.524,  0.   ,  0.524,  1.571,    nan], dtype=float32)

For complex-valued inputs:

>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.arcsin(3+4j)
Array(0.634+2.306j, dtype=complex64, weak_type=True)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.arcsinh(x, /)

Calculate element-wise inverse of hyperbolic sine of input.

JAX implementation of numpy.arcsinh.

The inverse of hyperbolic sine is defined by:

\[arcsinh(x) = \ln(x + \sqrt{1 + x^2})\]
Args:

x: input array or scalar.

Returns:

An array of same shape as x containing the inverse of hyperbolic sine of each element of x, promoting to inexact dtype.

Note:

  • jnp.arcsinh returns nan for values outside the range (-inf, inf).

  • jnp.arcsinh follows the branch cut convention of numpy.arcsinh for complex inputs.

See also:

  • jax.numpy.sinh(): Computes the element-wise hyperbolic sine of the input.

  • jax.numpy.arccosh(): Computes the element-wise inverse of hyperbolic cosine of the input.

  • jax.numpy.arctanh(): Computes the element-wise inverse of hyperbolic tangent of the input.

Examples:
>>> x = jnp.array([[-2, 3, 1],
...                [4, 9, -5]])
>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.arcsinh(x)
Array([[-1.444,  1.818,  0.881],
       [ 2.095,  2.893, -2.312]], dtype=float32)

For complex-valued inputs:

>>> x1 = jnp.array([4-3j, 2j])
>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.arcsinh(x1)
Array([2.306-0.634j, 1.317+1.571j], dtype=complex64)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.arctan(x, /)

Compute element-wise inverse of trigonometric tangent of input.

JAX implement of numpy.arctan.

Args:

x: input array or scalar.

Returns:

An array containing the inverse trigonometric tangent of each element x in radians in the range [-pi/2, pi/2], promoting to inexact dtype.

Note:

jnp.arctan follows the branch cut convention of numpy.arctan for complex inputs.

See also:

  • jax.numpy.tan(): Computes a trigonometric tangent of each element of input.

  • jax.numpy.arcsin() and jax.numpy.asin(): Computes the inverse of trigonometric sine of each element of input.

  • jax.numpy.arccos() and jax.numpy.atan(): Computes the inverse of trigonometric cosine of each element of input.

Examples:
>>> x = jnp.array([-jnp.inf, -20, -1, 0, 1, 20, jnp.inf])
>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.arctan(x)
Array([-1.571, -1.521, -0.785,  0.   ,  0.785,  1.521,  1.571], dtype=float32)

For complex-valued inputs:

>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.arctan(2+7j)
Array(1.532+0.133j, dtype=complex64, weak_type=True)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.arctan2(x1, x2, /)

Compute the arctangent of x1/x2, choosing the correct quadrant.

JAX implementation of numpy.arctan2()

Args:

x1: numerator array. x2: denomniator array; should be broadcast-compatible with x1.

Returns:

The elementwise arctangent of x1 / x2, tracking the correct quadrant.

See also:

  • jax.numpy.tan(): compute the tangent of an angle

  • jax.numpy.atan2(): the array API version of this function.

Examples:

Consider a sequence of angles in radians between 0 and \(2\pi\):

>>> theta = jnp.linspace(-jnp.pi, jnp.pi, 9)
>>> with jnp.printoptions(precision=2, suppress=True):
...   print(theta)
[-3.14 -2.36 -1.57 -0.79  0.    0.79  1.57  2.36  3.14]

These angles can equivalently be represented by (x, y) coordinates on a unit circle:

>>> x, y = jnp.cos(theta), jnp.sin(theta)

To reconstruct the input angle, we might be tempted to use the identity \(\tan(\theta) = y / x\), and compute \(\theta = \tan^{-1}(y/x)\). Unfortunately, this does not recover the input angle:

>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jnp.arctan(y / x))
[-0.    0.79  1.57 -0.79  0.    0.79  1.57 -0.79  0.  ]

The problem is that \(y/x\) contains some ambiguity: although \((y, x) = (-1, -1)\) and \((y, x) = (1, 1)\) represent different points in Cartesian space, in both cases \(y / x = 1\), and so the simple arctan approach loses information about which quadrant the angle lies in. arctan2() is built to address this:

>>> with jnp.printoptions(precision=2, suppress=True):
...  print(jnp.arctan2(y, x))
[ 3.14 -2.36 -1.57 -0.79  0.    0.79  1.57  2.36 -3.14]

The results match the input theta, except at the endpoints where \(+\pi\) and \(-\pi\) represent indistinguishable points on the unit circle. By convention, arctan2() always returns values between \(-\pi\) and \(+\pi\) inclusive.

Parameters:
Return type:

Array

quchip.declarative.qnp.arctanh(x, /)

Calculate element-wise inverse of hyperbolic tangent of input.

JAX implementation of numpy.arctanh.

The inverse of hyperbolic tangent is defined by:

\[arctanh(x) = \frac{1}{2} [\ln(1 + x) - \ln(1 - x)]\]
Args:

x: input array or scalar.

Returns:

An array of same shape as x containing the inverse of hyperbolic tangent of each element of x, promoting to inexact dtype.

Note:

  • jnp.arctanh returns nan for real-values outside the range [-1, 1].

  • jnp.arctanh follows the branch cut convention of numpy.arctanh for complex inputs.

See also:

  • jax.numpy.tanh(): Computes the element-wise hyperbolic tangent of the input.

  • jax.numpy.arcsinh(): Computes the element-wise inverse of hyperbolic sine of the input.

  • jax.numpy.arccosh(): Computes the element-wise inverse of hyperbolic cosine of the input.

Examples:
>>> x = jnp.array([-2, -1, -0.5, 0, 0.5, 1, 2])
>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.arctanh(x)
Array([   nan,   -inf, -0.549,  0.   ,  0.549,    inf,    nan], dtype=float32)

For complex-valued input:

>>> x1 = jnp.array([-2+0j, 3+0j, 4-1j])
>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.arctanh(x1)
Array([-0.549+1.571j,  0.347+1.571j,  0.239-1.509j], dtype=complex64)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.argmax(a, axis=None, out=None, keepdims=None)

Return the index of the maximum value of an array.

JAX implementation of numpy.argmax().

Args:

a: input array axis: optional integer specifying the axis along which to find the maximum

value. If axis is not specified, a will be flattened.

out: unused by JAX keepdims: if True, then return an array with the same number of dimensions

as a.

Returns:

an array containing the index of the maximum value along the specified axis.

See also:

  • jax.numpy.argmin(): return the index of the minimum value.

  • jax.numpy.nanargmax(): compute argmax while ignoring NaN values.

Examples:
>>> x = jnp.array([1, 3, 5, 4, 2])
>>> jnp.argmax(x)
Array(2, dtype=int32)

>>> x = jnp.array([[1, 3, 2],
...                [5, 4, 1]])
>>> jnp.argmax(x, axis=1)
Array([1, 0], dtype=int32)

>>> jnp.argmax(x, axis=1, keepdims=True)
Array([[1],
       [0]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.argmin(a, axis=None, out=None, keepdims=None)

Return the index of the minimum value of an array.

JAX implementation of numpy.argmin().

Args:

a: input array axis: optional integer specifying the axis along which to find the minimum

value. If axis is not specified, a will be flattened.

out: unused by JAX keepdims: if True, then return an array with the same number of dimensions

as a.

Returns:

an array containing the index of the minimum value along the specified axis.

See also:

  • jax.numpy.argmax(): return the index of the maximum value.

  • jax.numpy.nanargmin(): compute argmin while ignoring NaN values.

Examples:
>>> x = jnp.array([1, 3, 5, 4, 2])
>>> jnp.argmin(x)
Array(0, dtype=int32)

>>> x = jnp.array([[1, 3, 2],
...                [5, 4, 1]])
>>> jnp.argmin(x, axis=1)
Array([0, 2], dtype=int32)

>>> jnp.argmin(x, axis=1, keepdims=True)
Array([[0],
       [2]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.argpartition(a, kth, axis=-1)

Returns indices that partially sort an array.

JAX implementation of numpy.argpartition(). The JAX version differs from NumPy in the treatment of NaN entries: NaNs which have the negative bit set are sorted to the beginning of the array.

Args:

a: array to be partitioned. kth: static integer index about which to partition the array. axis: static integer axis along which to partition the array; default is -1.

Returns:

Indices which partition a at the kth value along axis. The entries before kth are indices of values smaller than take(a, kth, axis), and entries after kth are indices of values larger than take(a, kth, axis)

Note:

The JAX version requires the kth argument to be a static integer rather than a general array. This is implemented via two calls to jax.lax.top_k(). If you’re only accessing the top or bottom k values of the output, it may be more efficient to call jax.lax.top_k() directly.

See Also:

  • jax.numpy.partition(): direct partial sort

  • jax.numpy.argsort(): full indirect sort

  • jax.lax.top_k(): directly find the top k entries

  • jax.lax.approx_max_k(): compute the approximate top k entries

  • jax.lax.approx_min_k(): compute the approximate bottom k entries

Examples:
>>> x = jnp.array([6, 8, 4, 3, 1, 9, 7, 5, 2, 3])
>>> kth = 4
>>> idx = jnp.argpartition(x, kth)
>>> idx
Array([4, 8, 3, 9, 2, 0, 1, 5, 6, 7], dtype=int32)

The result is a sequence of indices that partially sort the input. All indices before kth are of values smaller than the pivot value, and all indices after kth are of values larger than the pivot value:

>>> x_partitioned = x[idx]
>>> smallest_values = x_partitioned[:kth]
>>> pivot_value = x_partitioned[kth]
>>> largest_values = x_partitioned[kth + 1:]
>>> print(smallest_values, pivot_value, largest_values)
[1 2 3 3] 4 [6 8 9 7 5]

Notice that among smallest_values and largest_values, the returned order is arbitrary and implementation-dependent.

Parameters:
Return type:

Array

quchip.declarative.qnp.argsort(a, axis=-1, *, kind=None, order=None, stable=True, descending=False)

Return indices that sort an array.

JAX implementation of numpy.argsort().

Args:

a: array to sort axis: integer axis along which to sort. Defaults to -1, i.e. the last

axis. If None, then a is flattened before being sorted.

stable: boolean specifying whether a stable sort should be used. Default=True. descending: boolean specifying whether to sort in descending order. Default=False. kind: deprecated; instead specify sort algorithm using stable=True or stable=False. order: not supported by JAX

Returns:

Array of indices that sort an array. Returned array will be of shape a.shape (if axis is an integer) or of shape (a.size,) (if axis is None).

Examples:

Simple 1-dimensional sort

>>> x = jnp.array([1, 3, 5, 4, 2, 1])
>>> indices = jnp.argsort(x)
>>> indices
Array([0, 5, 4, 1, 3, 2], dtype=int32)
>>> x[indices]
Array([1, 1, 2, 3, 4, 5], dtype=int32)

Sort along the last axis of an array:

>>> x = jnp.array([[2, 1, 3],
...                [6, 4, 3]])
>>> indices = jnp.argsort(x, axis=1)
>>> indices
Array([[1, 0, 2],
       [2, 1, 0]], dtype=int32)
>>> jnp.take_along_axis(x, indices, axis=1)
Array([[1, 2, 3],
       [3, 4, 6]], dtype=int32)
See also:

  • jax.numpy.sort(): return sorted values directly.

  • jax.numpy.lexsort(): lexicographical sort of multiple arrays.

  • jax.lax.sort(): lower-level function wrapping XLA’s Sort operator.

Parameters:
Return type:

Array

quchip.declarative.qnp.argwhere(a, *, size=None, fill_value=None)

Find the indices of nonzero array elements

JAX implementation of numpy.argwhere().

jnp.argwhere(x) is essentially equivalent to jnp.column_stack(jnp.nonzero(x)) with special handling for zero-dimensional (i.e. scalar) inputs.

Because the size of the output of argwhere is data-dependent, the function is not typically compatible with JIT. The JAX version adds the optional size argument, which specifies the size of the leading dimension of the output - it must be specified statically for jnp.argwhere to be compiled with non-static operands. See jax.numpy.nonzero() for a full discussion of size and its semantics.

Args:

a: array for which to find nonzero elements size: optional integer specifying statically the number of expected nonzero elements.

This must be specified in order to use argwhere within JAX transformations like jax.jit(). See jax.numpy.nonzero() for more information.

fill_value: optional array specifying the fill value when size is specified.

See jax.numpy.nonzero() for more information.

Returns:

a two-dimensional array of shape [size, x.ndim]. If size is not specified as an argument, it is equal to the number of nonzero elements in x.

See Also:

  • jax.numpy.where()

  • jax.numpy.nonzero()

Examples:

Two-dimensional array:

>>> x = jnp.array([[1, 0, 2],
...                [0, 3, 0]])
>>> jnp.argwhere(x)
Array([[0, 0],
       [0, 2],
       [1, 1]], dtype=int32)

Equivalent computation using jax.numpy.column_stack() and jax.numpy.nonzero():

>>> jnp.column_stack(jnp.nonzero(x))
Array([[0, 0],
       [0, 2],
       [1, 1]], dtype=int32)

Special case for zero-dimensional (i.e. scalar) inputs:

>>> jnp.argwhere(1)
Array([], shape=(1, 0), dtype=int32)
>>> jnp.argwhere(0)
Array([], shape=(0, 0), dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.around(a, decimals=0, out=None)

Alias of jax.numpy.round()

Parameters:
Return type:

Array

quchip.declarative.qnp.array(object, dtype=None, copy=True, order='K', ndmin=0, *, device=None)

Convert an object to a JAX array.

JAX implementation of numpy.array().

Args:
object: an object that is convertible to an array. This includes JAX

arrays, NumPy arrays, Python scalars, Python collections like lists and tuples, objects with an __array__ method, and objects supporting the Python buffer protocol.

dtype: optionally specify the dtype of the output array. If not

specified it will be inferred from the input.

copy: specify whether to force a copy of the input. Default: True. order: not implemented in JAX ndmin: integer specifying the minimum number of dimensions in the

output array.

device: optional Device or Sharding

to which the created array will be committed.

Returns:

A JAX array constructed from the input.

See also:

  • jax.numpy.asarray(): like array, but by default only copies when necessary.

  • jax.numpy.from_dlpack(): construct a JAX array from an object that implements the dlpack interface.

  • jax.numpy.frombuffer(): construct a JAX array from an object that implements the buffer interface.

Examples:

Constructing JAX arrays from Python scalars:

>>> jnp.array(True)
Array(True, dtype=bool)
>>> jnp.array(42)
Array(42, dtype=int32, weak_type=True)
>>> jnp.array(3.5)
Array(3.5, dtype=float32, weak_type=True)
>>> jnp.array(1 + 1j)
Array(1.+1.j, dtype=complex64, weak_type=True)

Constructing JAX arrays from Python collections:

>>> jnp.array([1, 2, 3])  # list of ints -> 1D array
Array([1, 2, 3], dtype=int32)
>>> jnp.array([(1, 2, 3), (4, 5, 6)])  # list of tuples of ints -> 2D array
Array([[1, 2, 3],
       [4, 5, 6]], dtype=int32)
>>> jnp.array(range(5))
Array([0, 1, 2, 3, 4], dtype=int32)

Constructing JAX arrays from NumPy arrays:

>>> jnp.array(np.linspace(0, 2, 5))
Array([0. , 0.5, 1. , 1.5, 2. ], dtype=float32)

Constructing a JAX array via the Python buffer interface, using Python’s built-in array module.

>>> from array import array
>>> pybuffer = array('i', [2, 3, 5, 7])
>>> jnp.array(pybuffer)
Array([2, 3, 5, 7], dtype=int32)
Parameters:

  • object (Any)

  • dtype (str | type[Any] | dtype | SupportsDType | None)

  • copy (bool)

  • order (str | None)

  • ndmin (int)

  • device (Device | Sharding | None)

Return type:

Array

quchip.declarative.qnp.array_equal(a1, a2, equal_nan=False)

Check if two arrays are element-wise equal.

JAX implementation of numpy.array_equal().

Args:

a1: first input array to compare. a2: second input array to compare. equal_nan: Boolean. If True, NaNs in a1 will be considered

equal to NaNs in a2. Default is False.

Returns:

Boolean scalar array indicating whether the input arrays are element-wise equal.

See Also:

  • jax.numpy.allclose()

  • jax.numpy.array_equiv()

Examples:
>>> jnp.array_equal(jnp.array([1, 2, 3]), jnp.array([1, 2, 3]))
Array(True, dtype=bool)
>>> jnp.array_equal(jnp.array([1, 2, 3]), jnp.array([1, 2]))
Array(False, dtype=bool)
>>> jnp.array_equal(jnp.array([1, 2, 3]), jnp.array([1, 2, 4]))
Array(False, dtype=bool)
>>> jnp.array_equal(jnp.array([1, 2, float('nan')]),
...                 jnp.array([1, 2, float('nan')]))
Array(False, dtype=bool)
>>> jnp.array_equal(jnp.array([1, 2, float('nan')]),
...                 jnp.array([1, 2, float('nan')]), equal_nan=True)
Array(True, dtype=bool)
Parameters:
Return type:

Array

quchip.declarative.qnp.array_equiv(a1, a2)

Check if two arrays are element-wise equal.

JAX implementation of numpy.array_equiv().

This function will return False if the input arrays cannot be broadcasted to the same shape.

Args:

a1: first input array to compare. a2: second input array to compare.

Returns:

Boolean scalar array indicating whether the input arrays are element-wise equal after broadcasting.

See Also:

  • jax.numpy.allclose()

  • jax.numpy.array_equal()

Examples:
>>> jnp.array_equiv(jnp.array([1, 2, 3]), jnp.array([1, 2, 3]))
Array(True, dtype=bool)
>>> jnp.array_equiv(jnp.array([1, 2, 3]), jnp.array([1, 2, 4]))
Array(False, dtype=bool)
>>> jnp.array_equiv(jnp.array([[1, 2, 3], [1, 2, 3]]),
...                 jnp.array([1, 2, 3]))
Array(True, dtype=bool)
Parameters:
Return type:

Array

quchip.declarative.qnp.array_repr(arr, max_line_width=None, precision=None, suppress_small=None)

Return the string representation of an array.

Parameters:

  • arr (ndarray) – Input array.

  • max_line_width (int, optional) – Inserts newlines if text is longer than max_line_width. Defaults to numpy.get_printoptions()['linewidth'].

  • precision (int, optional) – Floating point precision. Defaults to numpy.get_printoptions()['precision'].

  • suppress_small (bool, optional) – Represent numbers “very close” to zero as zero; default is False. Very close is defined by precision: if the precision is 8, e.g., numbers smaller (in absolute value) than 5e-9 are represented as zero. Defaults to numpy.get_printoptions()['suppress'].

Returns:

string – The string representation of an array.

Return type:

str

See also

array_str, array2string, set_printoptions

Examples

>>> import numpy as np
>>> np.array_repr(np.array([1,2]))
'array([1, 2])'
>>> np.array_repr(np.ma.array([0.]))
'MaskedArray([0.])'
>>> np.array_repr(np.array([], np.int32))
'array([], dtype=int32)'

>>> x = np.array([1e-6, 4e-7, 2, 3])
>>> np.array_repr(x, precision=6, suppress_small=True)
'array([0.000001,  0.      ,  2.      ,  3.      ])'
quchip.declarative.qnp.array_split(ary, indices_or_sections, axis=0)

Split an array into sub-arrays.

JAX implementation of numpy.array_split().

Refer to the documentation of jax.numpy.split() for details; array_split is equivalent to split, but allows integer indices_or_sections which does not evenly divide the split axis.

Examples:
>>> x = jnp.array([1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> chunks = jnp.array_split(x, 4)
>>> print(*chunks)
[1 2 3] [4 5] [6 7] [8 9]
See also:

  • jax.numpy.split(): split an array along any axis.

  • jax.numpy.vsplit(): split vertically, i.e. along axis=0

  • jax.numpy.hsplit(): split horizontally, i.e. along axis=1

  • jax.numpy.dsplit(): split depth-wise, i.e. along axis=2

Parameters:
Return type:

list[Array]

quchip.declarative.qnp.array_str(a, max_line_width=None, precision=None, suppress_small=None)

Return a string representation of the data in an array.

The data in the array is returned as a single string. This function is similar to array_repr, the difference being that array_repr also returns information on the kind of array and its data type.

Parameters:

  • a (ndarray) – Input array.

  • max_line_width (int, optional) – Inserts newlines if text is longer than max_line_width. Defaults to numpy.get_printoptions()['linewidth'].

  • precision (int, optional) – Floating point precision. Defaults to numpy.get_printoptions()['precision'].

  • suppress_small (bool, optional) – Represent numbers “very close” to zero as zero; default is False. Very close is defined by precision: if the precision is 8, e.g., numbers smaller (in absolute value) than 5e-9 are represented as zero. Defaults to numpy.get_printoptions()['suppress'].

See also

array2string, array_repr, set_printoptions

Examples

>>> import numpy as np
>>> np.array_str(np.arange(3))
'[0 1 2]'
quchip.declarative.qnp.asarray(a, dtype=None, order=None, *, copy=None, device=None)

Convert an object to a JAX array.

JAX implementation of numpy.asarray().

Args:
a: an object that is convertible to an array. This includes JAX

arrays, NumPy arrays, Python scalars, Python collections like lists and tuples, objects with an __array__ method, and objects supporting the Python buffer protocol.

dtype: optionally specify the dtype of the output array. If not

specified it will be inferred from the input.

order: not implemented in JAX copy: optional boolean specifying the copy mode. If True, then always

return a copy. If False, then error if a copy is necessary. Default is None, which will only copy when necessary.

device: optional Device or Sharding

to which the created array will be committed.

Returns:

A JAX array constructed from the input.

See also:

  • jax.numpy.array(): like asarray, but defaults to copy=True.

  • jax.numpy.from_dlpack(): construct a JAX array from an object that implements the dlpack interface.

  • jax.numpy.frombuffer(): construct a JAX array from an object that implements the buffer interface.

Examples:

Constructing JAX arrays from Python scalars:

>>> jnp.asarray(True)
Array(True, dtype=bool)
>>> jnp.asarray(42)
Array(42, dtype=int32, weak_type=True)
>>> jnp.asarray(3.5)
Array(3.5, dtype=float32, weak_type=True)
>>> jnp.asarray(1 + 1j)
Array(1.+1.j, dtype=complex64, weak_type=True)

Constructing JAX arrays from Python collections:

>>> jnp.asarray([1, 2, 3])  # list of ints -> 1D array
Array([1, 2, 3], dtype=int32)
>>> jnp.asarray([(1, 2, 3), (4, 5, 6)])  # list of tuples of ints -> 2D array
Array([[1, 2, 3],
       [4, 5, 6]], dtype=int32)
>>> jnp.asarray(range(5))
Array([0, 1, 2, 3, 4], dtype=int32)

Constructing JAX arrays from NumPy arrays:

>>> jnp.asarray(np.linspace(0, 2, 5))
Array([0. , 0.5, 1. , 1.5, 2. ], dtype=float32)

Constructing a JAX array via the Python buffer interface, using Python’s built-in array module.

>>> from array import array
>>> pybuffer = array('i', [2, 3, 5, 7])
>>> jnp.asarray(pybuffer)
Array([2, 3, 5, 7], dtype=int32)
Parameters:

  • a (Any)

  • dtype (str | type[Any] | dtype | SupportsDType | None)

  • order (str | None)

  • copy (bool | None)

  • device (Device | Sharding | None)

Return type:

Array

quchip.declarative.qnp.asin(x, /)

Alias of jax.numpy.arcsin()

Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.asinh(x, /)

Alias of jax.numpy.arcsinh()

Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.astype(x, dtype, /, *, copy=False, device=None)

Convert an array to a specified dtype.

JAX implementation of numpy.astype().

This is implemented via jax.lax.convert_element_type(), which may have slightly different behavior than numpy.astype() in some cases. In particular, the details of float-to-int and int-to-float casts are implementation dependent.

Args:

x: input array to convert dtype: output dtype copy: if True, then always return a copy. If False (default) then only

return a copy if necessary.

device: optionally specify the device to which the output will be committed.

Returns:

An array with the same shape as x, containing values of the specified dtype.

See Also:

  • jax.lax.convert_element_type(): lower-level function for XLA-style dtype conversions.

Examples:
>>> x = jnp.array([0, 1, 2, 3])
>>> x
Array([0, 1, 2, 3], dtype=int32)
>>> x.astype('float32')
Array([0.0, 1.0, 2.0, 3.0], dtype=float32)

>>> y = jnp.array([0.0, 0.5, 1.0])
>>> y.astype(int)  # truncates fractional values
Array([0, 0, 1], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.atan(x, /)

Alias of jax.numpy.arctan()

Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.atan2(x1, x2, /)

Alias of jax.numpy.arctan2()

Parameters:
Return type:

Array

quchip.declarative.qnp.atanh(x, /)

Alias of jax.numpy.arctanh()

Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.atleast_1d(*arys)

Convert inputs to arrays with at least 1 dimension.

JAX implementation of numpy.atleast_1d().

Args:

zero or more arraylike arguments.

Returns:

an array or list of arrays corresponding to the input values. Arrays of shape () are converted to shape (1,), and arrays with other shapes are returned unchanged.

See also:

  • jax.numpy.asarray()

  • jax.numpy.atleast_2d()

  • jax.numpy.atleast_3d()

Examples:

Scalar arguments are converted to 1D, length-1 arrays:

>>> x = jnp.float32(1.0)
>>> jnp.atleast_1d(x)
Array([1.], dtype=float32)

Higher dimensional inputs are returned unchanged:

>>> y = jnp.arange(4)
>>> jnp.atleast_1d(y)
Array([0, 1, 2, 3], dtype=int32)

Multiple arguments can be passed to the function at once, in which case a list of results is returned:

>>> jnp.atleast_1d(x, y)
[Array([1.], dtype=float32), Array([0, 1, 2, 3], dtype=int32)]
Parameters:

arys (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array | list[Array]

quchip.declarative.qnp.atleast_2d(*arys)

Convert inputs to arrays with at least 2 dimensions.

JAX implementation of numpy.atleast_2d().

Args:

zero or more arraylike arguments.

Returns:

an array or list of arrays corresponding to the input values. Arrays of shape () are converted to shape (1, 1), 1D arrays of shape (N,) are converted to shape (1, N), and arrays of all other shapes are returned unchanged.

See also:

  • jax.numpy.asarray()

  • jax.numpy.atleast_1d()

  • jax.numpy.atleast_3d()

Examples:

Scalar arguments are converted to 2D, size-1 arrays:

>>> x = jnp.float32(1.0)
>>> jnp.atleast_2d(x)
Array([[1.]], dtype=float32)

One-dimensional arguments have a unit dimension prepended to the shape:

>>> y = jnp.arange(4)
>>> jnp.atleast_2d(y)
Array([[0, 1, 2, 3]], dtype=int32)

Higher dimensional inputs are returned unchanged:

>>> z = jnp.ones((2, 3))
>>> jnp.atleast_2d(z)
Array([[1., 1., 1.],
       [1., 1., 1.]], dtype=float32)

Multiple arguments can be passed to the function at once, in which case a list of results is returned:

>>> jnp.atleast_2d(x, y)
[Array([[1.]], dtype=float32), Array([[0, 1, 2, 3]], dtype=int32)]
Parameters:

arys (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array | list[Array]

quchip.declarative.qnp.atleast_3d(*arys)

Convert inputs to arrays with at least 3 dimensions.

JAX implementation of numpy.atleast_3d().

Args:

zero or more arraylike arguments.

Returns:

an array or list of arrays corresponding to the input values. Arrays of shape () are converted to shape (1, 1, 1), 1D arrays of shape (N,) are converted to shape (1, N, 1), 2D arrays of shape (M, N) are converted to shape (M, N, 1), and arrays of all other shapes are returned unchanged.

See also:

  • jax.numpy.asarray()

  • jax.numpy.atleast_1d()

  • jax.numpy.atleast_2d()

Examples:

Scalar arguments are converted to 3D, size-1 arrays:

>>> x = jnp.float32(1.0)
>>> jnp.atleast_3d(x)
Array([[[1.]]], dtype=float32)

1D arrays have a unit dimension prepended and appended:

>>> y = jnp.arange(4)
>>> jnp.atleast_3d(y).shape
(1, 4, 1)

2D arrays have a unit dimension appended:

>>> z = jnp.ones((2, 3))
>>> jnp.atleast_3d(z).shape
(2, 3, 1)

Multiple arguments can be passed to the function at once, in which case a list of results is returned:

>>> x3, y3 = jnp.atleast_3d(x, y)
>>> print(x3)
[[[1.]]]
>>> print(y3)
[[[0]
  [1]
  [2]
  [3]]]
Parameters:

arys (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array | list[Array]

quchip.declarative.qnp.average(a, axis=None, weights=None, returned=False, keepdims=False)

Compute the weighed average.

JAX Implementation of numpy.average().

Args:

a: array to be averaged axis: an optional integer or sequence of integers specifying the axis along which

the mean to be computed. If not specified, mean is computed along all the axes.

weights: an optional array of weights for a weighted average. Must be

broadcast-compatible with a.

returned: If False (default) then return only the average. If True then return both

the average and the normalization factor (i.e. the sum of weights).

keepdims: If True, reduced axes are left in the result with size 1. If False (default)

then reduced axes are squeezed out.

Returns:

An array average or tuple of arrays (average, normalization) if returned is True.

See also:

  • jax.numpy.mean(): unweighted mean.

Examples:

Simple average:

>>> x = jnp.array([1, 2, 3, 2, 4])
>>> jnp.average(x)
Array(2.4, dtype=float32)

Weighted average:

>>> weights = jnp.array([2, 1, 3, 2, 2])
>>> jnp.average(x, weights=weights)
Array(2.5, dtype=float32)

Use returned=True to optionally return the normalization, i.e. the sum of weights:

>>> jnp.average(x, returned=True)
(Array(2.4, dtype=float32), Array(5., dtype=float32))
>>> jnp.average(x, weights=weights, returned=True)
(Array(2.5, dtype=float32), Array(10., dtype=float32))

Weighted average along a specified axis:

>>> x = jnp.array([[8, 2, 7],
...                [3, 6, 4]])
>>> weights = jnp.array([1, 2, 3])
>>> jnp.average(x, weights=weights, axis=1)
Array([5.5, 4.5], dtype=float32)
Parameters:
Return type:

Array | tuple[Array, Array]

quchip.declarative.qnp.bartlett(M)

Return a Bartlett window of size M.

JAX implementation of numpy.bartlett().

Args:

M: The window size.

Returns:

An array of size M containing the Bartlett window.

Examples:
>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jnp.bartlett(4))
[0.   0.67 0.67 0.  ]
See also:

  • jax.numpy.blackman(): return a Blackman window of size M.

  • jax.numpy.hamming(): return a Hamming window of size M.

  • jax.numpy.hanning(): return a Hanning window of size M.

  • jax.numpy.kaiser(): return a Kaiser window of size M.

Parameters:

M (int)

Return type:

Array

class quchip.declarative.qnp.bfloat16(x)

Bases: object

A JAX scalar constructor of type bfloat16.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype(bfloat16)
quchip.declarative.qnp.bincount(x, weights=None, minlength=0, *, length=None)

Count the number of occurrences of each value in an integer array.

JAX implementation of numpy.bincount().

For an array of non-negative integers x, this function returns an array counts of size x.max() + 1, such that counts[i] contains the number of occurrences of the value i in x.

The JAX version has a few differences from the NumPy version:

  • In NumPy, passing an array x with negative entries will result in an error. In JAX, negative values are clipped to zero.

  • JAX adds an optional length parameter which can be used to statically specify the length of the output array so that this function can be used with transformations like jax.jit(). In this case, items larger than length + 1 will be dropped.

Args:

x : 1-dimensional array of non-negative integers weights: optional array of weights associated with x. If not specified, the

weight for each entry will be 1.

minlength: the minimum length of the output counts array. length: the length of the output counts array. Must be specified statically for

bincount to be used with jax.jit() and other JAX transformations.

Returns:

An array of counts or summed weights reflecting the number of occurrences of values in x.

See Also:

  • jax.numpy.histogram()

  • jax.numpy.digitize()

  • jax.numpy.unique_counts()

Examples:

Basic bincount:

>>> x = jnp.array([1, 1, 2, 3, 3, 3])
>>> jnp.bincount(x)
Array([0, 2, 1, 3], dtype=int32)

Weighted bincount:

>>> weights = jnp.array([1, 2, 3, 4, 5, 6])
>>> jnp.bincount(x, weights)
Array([ 0,  3,  3, 15], dtype=int32)

Specifying a static length makes this jit-compatible:

>>> jit_bincount = jax.jit(jnp.bincount, static_argnames=['length'])
>>> jit_bincount(x, length=5)
Array([0, 2, 1, 3, 0], dtype=int32)

Any negative numbers are clipped to the first bin, and numbers beyond the specified length are dropped:

>>> x = jnp.array([-1, -1, 1, 3, 10])
>>> jnp.bincount(x, length=5)
Array([2, 1, 0, 1, 0], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.bitwise_count(x, /)

Counts the number of 1 bits in the binary representation of the absolute value of each element of x.

JAX implementation of numpy.bitwise_count.

Args:

x: Input array, only accepts integer subtypes

Returns:

An array-like object containing the binary 1 bit counts of the absolute value of each element in x, with the same shape as x of dtype uint8.

Examples:
>>> x1 = jnp.array([64, 32, 31, 20])
>>> # 64 = 0b1000000, 32 = 0b100000, 31 = 0b11111, 20 = 0b10100
>>> jnp.bitwise_count(x1)
Array([1, 1, 5, 2], dtype=uint8)

>>> x2 = jnp.array([-16, -7, 7])
>>> # |-16| = 0b10000, |-7| = 0b111, 7 = 0b111
>>> jnp.bitwise_count(x2)
Array([1, 3, 3], dtype=uint8)

>>> x3 = jnp.array([[2, -7],[-9, 7]])
>>> # 2 = 0b10, |-7| = 0b111, |-9| = 0b1001, 7 = 0b111
>>> jnp.bitwise_count(x3)
Array([[1, 3],
       [2, 3]], dtype=uint8)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.bitwise_invert(x, /)

Alias of jax.numpy.invert().

Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.bitwise_left_shift(x, y, /)

Alias of jax.numpy.left_shift().

Parameters:
Return type:

Array

quchip.declarative.qnp.bitwise_not(x, /)

Alias of jax.numpy.invert().

Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.bitwise_right_shift(x1, x2, /)

Alias of jax.numpy.right_shift().

Parameters:
Return type:

Array

quchip.declarative.qnp.blackman(M)

Return a Blackman window of size M.

JAX implementation of numpy.blackman().

Args:

M: The window size.

Returns:

An array of size M containing the Blackman window.

Examples:
>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jnp.blackman(4))
[-0.    0.63  0.63 -0.  ]
See also:

  • jax.numpy.bartlett(): return a Bartlett window of size M.

  • jax.numpy.hamming(): return a Hamming window of size M.

  • jax.numpy.hanning(): return a Hanning window of size M.

  • jax.numpy.kaiser(): return a Kaiser window of size M.

Parameters:

M (int)

Return type:

Array

quchip.declarative.qnp.block(arrays)

Create an array from a list of blocks.

JAX implementation of numpy.block().

Args:
arrays: an array, or nested list of arrays which will be concatenated

together to form the final array.

Returns:

a single array constructed from the inputs.

See also:
Examples:

consider these blocks:

>>> zeros = jnp.zeros((2, 2))
>>> ones = jnp.ones((2, 2))
>>> twos = jnp.full((2, 2), 2)
>>> threes = jnp.full((2, 2), 3)

Passing a single array to block() returns the array:

>>> jnp.block(zeros)
Array([[0., 0.],
       [0., 0.]], dtype=float32)

Passing a simple list of arrays concatenates them along the last axis:

>>> jnp.block([zeros, ones])
Array([[0., 0., 1., 1.],
       [0., 0., 1., 1.]], dtype=float32)

Passing a doubly-nested list of arrays concatenates the inner list along the last axis, and the outer list along the second-to-last axis:

>>> jnp.block([[zeros, ones],
...            [twos, threes]])
Array([[0., 0., 1., 1.],
       [0., 0., 1., 1.],
       [2., 2., 3., 3.],
       [2., 2., 3., 3.]], dtype=float32)

Note that blocks need not align in all dimensions, though the size along the axis of concatenation must match. For example, this is valid because after the inner, horizontal concatenation, the resulting blocks have a valid shape for the outer, vertical concatenation.

>>> a = jnp.zeros((2, 1))
>>> b = jnp.ones((2, 3))
>>> c = jnp.full((1, 2), 2)
>>> d = jnp.full((1, 2), 3)
>>> jnp.block([[a, b], [c, d]])
Array([[0., 1., 1., 1.],
       [0., 1., 1., 1.],
       [2., 2., 3., 3.]], dtype=float32)

Note also that this logic generalizes to blocks in 3 or more dimensions. Here’s a 3-dimensional block-wise array:

>>> x = jnp.arange(6).reshape((1, 2, 3))
>>> blocks = [[[x for i in range(3)] for j in range(4)] for k in range(5)]
>>> jnp.block(blocks).shape
(5, 8, 9)
Parameters:

arrays (Array | ndarray | bool | number | bool | int | float | complex | list[Array | ndarray | bool | number | bool | int | float | complex])

Return type:

Array

class quchip.declarative.qnp.bool(x)

Bases: object

A JAX scalar constructor of type bool.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype('bool')
quchip.declarative.qnp.bool_

alias of bool

quchip.declarative.qnp.broadcast_arrays(*args)

Broadcast arrays to a common shape.

JAX implementation of numpy.broadcast_arrays(). JAX uses NumPy-style broadcasting rules, which you can read more about at NumPy broadcasting.

Args:

args: zero or more array-like objects to be broadcasted.

Returns:

a list of arrays containing broadcasted copies of the inputs.

See also:

  • jax.numpy.broadcast_shapes(): broadcast input shapes to a common shape.

  • jax.numpy.broadcast_to(): broadcast an array to a specified shape.

Examples:

>>> x = jnp.arange(3)
>>> y = jnp.int32(1)
>>> jnp.broadcast_arrays(x, y)
[Array([0, 1, 2], dtype=int32), Array([1, 1, 1], dtype=int32)]

>>> x = jnp.array([[1, 2, 3]])
>>> y = jnp.array([[10],
...                [20]])
>>> x2, y2 = jnp.broadcast_arrays(x, y)
>>> x2
Array([[1, 2, 3],
       [1, 2, 3]], dtype=int32)
>>> y2
Array([[10, 10, 10],
       [20, 20, 20]], dtype=int32)
Parameters:

args (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

list[Array]

quchip.declarative.qnp.broadcast_shapes(*shapes)

Broadcast input shapes to a common output shape.

JAX implementation of numpy.broadcast_shapes(). JAX uses NumPy-style broadcasting rules, which you can read more about at NumPy broadcasting.

Args:

shapes: 0 or more shapes specified as sequences of integers

Returns:

The broadcasted shape as a tuple of integers.

See Also:

  • jax.numpy.broadcast_arrays(): broadcast arrays to a common shape.

  • jax.numpy.broadcast_to(): broadcast an array to a specified shape.

Examples:

Some compatible shapes:

>>> jnp.broadcast_shapes((1,), (4,))
(4,)
>>> jnp.broadcast_shapes((3, 1), (4,))
(3, 4)
>>> jnp.broadcast_shapes((3, 1), (1, 4), (5, 1, 1))
(5, 3, 4)

Incompatible shapes:

>>> jnp.broadcast_shapes((3, 1), (4, 1))
Traceback (most recent call last):
ValueError: Incompatible shapes for broadcasting: shapes=[(3, 1), (4, 1)]
quchip.declarative.qnp.broadcast_to(array, shape, *, out_sharding=None)

Broadcast an array to a specified shape.

JAX implementation of numpy.broadcast_to(). JAX uses NumPy-style broadcasting rules, which you can read more about at NumPy broadcasting.

Args:

array: array to be broadcast. shape: shape to which the array will be broadcast.

Returns:

a copy of array broadcast to the specified shape.

See also:

  • jax.numpy.broadcast_arrays(): broadcast arrays to a common shape.

  • jax.numpy.broadcast_shapes(): broadcast input shapes to a common shape.

Examples:
>>> x = jnp.int32(1)
>>> jnp.broadcast_to(x, (1, 4))
Array([[1, 1, 1, 1]], dtype=int32)

>>> x = jnp.array([1, 2, 3])
>>> jnp.broadcast_to(x, (2, 3))
Array([[1, 2, 3],
       [1, 2, 3]], dtype=int32)

>>> x = jnp.array([[2], [4]])
>>> jnp.broadcast_to(x, (2, 4))
Array([[2, 2, 2, 2],
       [4, 4, 4, 4]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.can_cast(from_, to, casting='safe')

Returns True if cast between data types can occur according to the casting rule.

Parameters:

  • from (dtype, dtype specifier, NumPy scalar, or array) – Data type, NumPy scalar, or array to cast from.

  • to (dtype or dtype specifier) – Data type to cast to.

  • casting ({'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional) –

    Controls what kind of data casting may occur.

    • ’no’ means the data types should not be cast at all.

    • ’equiv’ means only byte-order changes are allowed.

    • ’safe’ means only casts which can preserve values are allowed.

    • ’same_kind’ means only safe casts or casts within a kind, like float64 to float32, are allowed.

    • ’unsafe’ means any data conversions may be done.

Returns:

out – True if cast can occur according to the casting rule.

Return type:

bool

Notes

Changed in version 2.0: This function does not support Python scalars anymore and does not apply any value-based logic for 0-D arrays and NumPy scalars.

See also

dtype, result_type

Examples

Basic examples

>>> import numpy as np
>>> np.can_cast(np.int32, np.int64)
True
>>> np.can_cast(np.float64, complex)
True
>>> np.can_cast(complex, float)
False

>>> np.can_cast('i8', 'f8')
True
>>> np.can_cast('i8', 'f4')
False
>>> np.can_cast('i4', 'S4')
False
quchip.declarative.qnp.cbrt(x, /)

Calculates element-wise cube root of the input array.

JAX implementation of numpy.cbrt.

Args:

x: input array or scalar. complex dtypes are not supported.

Returns:

An array containing the cube root of the elements of x.

See also:

  • jax.numpy.sqrt(): Calculates the element-wise non-negative square root of the input.

  • jax.numpy.square(): Calculates the element-wise square of the input.

Examples:
>>> x = jnp.array([[216, 125, 64],
...                [-27, -8, -1]])
>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.cbrt(x)
Array([[ 6.,  5.,  4.],
       [-3., -2., -1.]], dtype=float32)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.cdouble

alias of complex128

quchip.declarative.qnp.ceil(x, /)

Round input to the nearest integer upwards.

JAX implementation of numpy.ceil.

Args:

x: input array or scalar. Must not have complex dtype.

Returns:

An array with same shape and dtype as x containing the values rounded to the nearest integer that is greater than or equal to the value itself.

See also:

  • jax.numpy.fix(): Rounds the input to the nearest integer towards zero.

  • jax.numpy.trunc(): Rounds the input to the nearest integer towards zero.

  • jax.numpy.floor(): Rounds the input down to the nearest integer.

Examples:
>>> key = jax.random.key(1)
>>> x = jax.random.uniform(key, (3, 3), minval=-5, maxval=5)
>>> with jnp.printoptions(precision=2, suppress=True):
...     print(x)
[[-0.61  0.34 -0.54]
 [-0.62  3.97  0.59]
 [ 4.84  3.42 -1.14]]
>>> jnp.ceil(x)
Array([[-0.,  1., -0.],
       [-0.,  4.,  1.],
       [ 5.,  4., -1.]], dtype=float32)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

class quchip.declarative.qnp.character

Bases: flexible

Abstract base class of all character string scalar types.

quchip.declarative.qnp.choose(a, choices, out=None, mode='raise')

Construct an array by stacking slices of choice arrays.

JAX implementation of numpy.choose().

The semantics of this function can be confusing, but in the simplest case where a is a one-dimensional array, choices is a two-dimensional array, and all entries of a are in-bounds (i.e. 0 <= a_i < len(choices)), then the function is equivalent to the following:

def choose(a, choices):
  return jnp.array([choices[a_i, i] for i, a_i in enumerate(a)])

In the more general case, a may have any number of dimensions and choices may be an arbitrary sequence of broadcast-compatible arrays. In this case, again for in-bound indices, the logic is equivalent to:

def choose(a, choices):
  a, *choices = jnp.broadcast_arrays(a, *choices)
  choices = jnp.array(choices)
  return jnp.array([choices[a[idx], *idx] for idx in np.ndindex(a.shape)])

The only additional complexity comes from the mode argument, which controls the behavior for out-of-bound indices in a as described below.

Args:

a: an N-dimensional array of integer indices. choices: an array or sequence of arrays. All arrays in the sequence must be

mutually broadcast compatible with a.

out: unused by JAX mode: specify the out-of-bounds indexing mode; one of 'raise' (default),

'wrap', or 'clip'. Note that the default mode of 'raise' is not compatible with JAX transformations.

Returns:

an array containing stacked slices from choices at the indices specified by a. The shape of the result is broadcast_shapes(a.shape, *(c.shape for c in choices)).

See also:

  • jax.lax.switch(): choose between N functions based on an index.

Examples:

Here is the simplest case of a 1D index array with a 2D choice array, in which case this chooses the indexed value from each column:

>>> choices = jnp.array([[ 1,  2,  3,  4],
...                      [ 5,  6,  7,  8],
...                      [ 9, 10, 11, 12]])
>>> a = jnp.array([2, 0, 1, 0])
>>> jnp.choose(a, choices)
Array([9, 2, 7, 4], dtype=int32)

The mode argument specifies what to do with out-of-bound indices; options are to either wrap or clip:

>>> a2 = jnp.array([2, 0, 1, 4])  # last index out-of-bound
>>> jnp.choose(a2, choices, mode='clip')
Array([ 9,  2,  7, 12], dtype=int32)
>>> jnp.choose(a2, choices, mode='wrap')
Array([9, 2, 7, 8], dtype=int32)

In the more general case, choices may be a sequence of array-like objects with any broadcast-compatible shapes.

>>> choice_1 = jnp.array([1, 2, 3, 4])
>>> choice_2 = 99
>>> choice_3 = jnp.array([[10],
...                       [20],
...                       [30]])
>>> a = jnp.array([[0, 1, 2, 0],
...                [1, 2, 0, 1],
...                [2, 0, 1, 2]])
>>> jnp.choose(a, [choice_1, choice_2, choice_3], mode='wrap')
Array([[ 1, 99, 10,  4],
       [99, 20,  3, 99],
       [30,  2, 99, 30]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.clip(arr=None, /, min=None, max=None, *, a=Deprecated, a_min=Deprecated, a_max=Deprecated)

Clip array values to a specified range.

JAX implementation of numpy.clip().

Args:

arr: N-dimensional array to be clipped. min: optional minimum value of the clipped range; if None (default) then

result will not be clipped to any minimum value. If specified, it should be broadcast-compatible with arr and max.

max: optional maximum value of the clipped range; if None (default) then

result will not be clipped to any maximum value. If specified, it should be broadcast-compatible with arr and min.

a: deprecated alias of the arr argument. Will result in a

DeprecationWarning if used.

a_min: deprecated alias of the min argument. Will result in a

DeprecationWarning if used.

a_max: deprecated alias of the max argument. Will result in a

DeprecationWarning if used.

Returns:

An array containing values from arr, with values smaller than min set to min, and values larger than max set to max. Wherever min is larger than max, the value of max is returned.

See also:

  • jax.numpy.minimum(): Compute the element-wise minimum value of two arrays.

  • jax.numpy.maximum(): Compute the element-wise maximum value of two arrays.

Examples:
>>> arr = jnp.array([0, 1, 2, 3, 4, 5, 6, 7])
>>> jnp.clip(arr, 2, 5)
Array([2, 2, 2, 3, 4, 5, 5, 5], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.column_stack(tup)

Stack arrays column-wise.

JAX implementation of numpy.column_stack().

For arrays of two or more dimensions, this is equivalent to jax.numpy.concatenate() with axis=1.

Args:
tup: a sequence of arrays to stack; each must have the same leading dimension.

Input arrays will be promoted to at least rank 2. If a single array is given it will be treated equivalently to tup = unstack(tup), but the implementation will avoid explicit unstacking.

dtype: optional dtype of the resulting array. If not specified, the dtype

will be determined via type promotion rules described in type-promotion.

Returns:

the stacked result.

See also:

  • jax.numpy.stack(): stack along arbitrary axes

  • jax.numpy.concatenate(): concatenation along existing axes.

  • jax.numpy.vstack(): stack vertically, i.e. along axis 0.

  • jax.numpy.hstack(): stack horizontally, i.e. along axis 1.

  • jax.numpy.hstack(): stack depth=wise, i.e. along axis 2.

Examples:

Scalar values:

>>> jnp.column_stack([1, 2, 3])
Array([[1, 2, 3]], dtype=int32, weak_type=True)

1D arrays:

>>> x = jnp.arange(3)
>>> y = jnp.ones(3)
>>> jnp.column_stack([x, y])
Array([[0., 1.],
       [1., 1.],
       [2., 1.]], dtype=float32)

2D arrays:

>>> x = x.reshape(3, 1)
>>> y = y.reshape(3, 1)
>>> jnp.column_stack([x, y])
Array([[0., 1.],
       [1., 1.],
       [2., 1.]], dtype=float32)
Parameters:

tup (ndarray | Array | Sequence[Array | ndarray | bool | number | bool | int | float | complex])

Return type:

Array

class quchip.declarative.qnp.complex128(x)

Bases: object

A JAX scalar constructor of type complex128.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype('complex128')
class quchip.declarative.qnp.complex64(x)

Bases: object

A JAX scalar constructor of type complex64.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype('complex64')
quchip.declarative.qnp.complex_

alias of complex128

class quchip.declarative.qnp.complexfloating

Bases: inexact

Abstract base class of all complex number scalar types that are made up of floating-point numbers.

quchip.declarative.qnp.compress(condition, a, axis=None, *, size=None, fill_value=0, out=None)

Compress an array along a given axis using a boolean condition.

JAX implementation of numpy.compress().

Args:

condition: 1-dimensional array of conditions. Will be converted to boolean. a: N-dimensional array of values. axis: axis along which to compress. If None (default) then a will be

flattened, and axis will be set to 0.

size: optional static size for output. Must be specified in order for compress

to be compatible with JAX transformations like jit() or vmap().

fill_value: if size is specified, fill padded entries with this value (default: 0). out: not implemented by JAX.

Returns:

An array of dimension a.ndim, compressed along the specified axis.

See also:

  • jax.numpy.extract(): 1D version of compress.

  • jax.Array.compress(): equivalent functionality as an array method.

Notes:

This function does not require strict shape agreement between condition and a. If condition.size > a.shape[axis], then condition will be truncated, and if a.shape[axis] > condition.size, then a will be truncated.

Examples:

Compressing along the rows of a 2D array:

>>> a = jnp.array([[1,  2,  3,  4],
...                [5,  6,  7,  8],
...                [9,  10, 11, 12]])
>>> condition = jnp.array([True, False, True])
>>> jnp.compress(condition, a, axis=0)
Array([[ 1,  2,  3,  4],
       [ 9, 10, 11, 12]], dtype=int32)

For convenience, you can equivalently use the compress() method of JAX arrays:

>>> a.compress(condition, axis=0)
Array([[ 1,  2,  3,  4],
       [ 9, 10, 11, 12]], dtype=int32)

Note that the condition need not match the shape of the specified axis; here we compress the columns with the length-3 condition. Values beyond the size of the condition are ignored:

>>> jnp.compress(condition, a, axis=1)
Array([[ 1,  3],
       [ 5,  7],
       [ 9, 11]], dtype=int32)

The optional size argument lets you specify a static output size so that the output is statically-shaped, and so this function can be used with transformations like jit() and vmap():

>>> f = lambda c, a: jnp.extract(c, a, size=len(a), fill_value=0)
>>> mask = (a % 3 == 0)
>>> jax.vmap(f)(mask, a)
Array([[ 3,  0,  0,  0],
       [ 6,  0,  0,  0],
       [ 9, 12,  0,  0]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.concat(arrays, /, *, axis=0)

Join arrays along an existing axis.

JAX implementation of array_api.concat().

Args:
arrays: a sequence of arrays to concatenate; each must have the same shape

except along the specified axis. If a single array is given it will be treated equivalently to arrays = unstack(arrays), but the implementation will avoid explicit unstacking.

axis: specify the axis along which to concatenate.

Returns:

the concatenated result.

See also:

  • jax.lax.concatenate(): XLA concatenation API.

  • jax.numpy.concatenate(): NumPy version of this function.

  • jax.numpy.stack(): concatenate arrays along a new axis.

Examples:

One-dimensional concatenation:

>>> x = jnp.arange(3)
>>> y = jnp.zeros(3, dtype=int)
>>> jnp.concat([x, y])
Array([0, 1, 2, 0, 0, 0], dtype=int32)

Two-dimensional concatenation:

>>> x = jnp.ones((2, 3))
>>> y = jnp.zeros((2, 1))
>>> jnp.concat([x, y], axis=1)
Array([[1., 1., 1., 0.],
       [1., 1., 1., 0.]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.concatenate(arrays, axis=0, dtype=None)

Join arrays along an existing axis.

JAX implementation of numpy.concatenate().

Args:
arrays: a sequence of arrays to concatenate; each must have the same shape

except along the specified axis. If a single array is given it will be treated equivalently to arrays = unstack(arrays), but the implementation will avoid explicit unstacking.

axis: specify the axis along which to concatenate. dtype: optional dtype of the resulting array. If not specified, the dtype

will be determined via type promotion rules described in type-promotion.

Returns:

the concatenated result.

See also:

  • jax.lax.concatenate(): XLA concatenation API.

  • jax.numpy.concat(): Array API version of this function.

  • jax.numpy.stack(): concatenate arrays along a new axis.

Examples:

One-dimensional concatenation:

>>> x = jnp.arange(3)
>>> y = jnp.zeros(3, dtype=int)
>>> jnp.concatenate([x, y])
Array([0, 1, 2, 0, 0, 0], dtype=int32)

Two-dimensional concatenation:

>>> x = jnp.ones((2, 3))
>>> y = jnp.zeros((2, 1))
>>> jnp.concatenate([x, y], axis=1)
Array([[1., 1., 1., 0.],
       [1., 1., 1., 0.]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.conj(x, /)

Alias of jax.numpy.conjugate()

Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.conjugate(x, /)

Return element-wise complex-conjugate of the input.

JAX implementation of numpy.conjugate.

Args:

x: inpuat array or scalar.

Returns:

An array containing the complex-conjugate of x.

See also:

  • jax.numpy.real(): Returns the element-wise real part of the complex argument.

  • jax.numpy.imag(): Returns the element-wise imaginary part of the complex argument.

Examples:
>>> jnp.conjugate(3)
Array(3, dtype=int32, weak_type=True)
>>> x = jnp.array([2-1j, 3+5j, 7])
>>> jnp.conjugate(x)
Array([2.+1.j, 3.-5.j, 7.-0.j], dtype=complex64)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.convolve(a, v, mode='full', *, precision=None, preferred_element_type=None)

Convolution of two one dimensional arrays.

JAX implementation of numpy.convolve().

Convolution of one dimensional arrays is defined as:

\[c_k = \sum_j a_{k - j} v_j\]
Args:

a: left-hand input to the convolution. Must have a.ndim == 1. v: right-hand input to the convolution. Must have v.ndim == 1. mode: controls the size of the output. Available operations are:

  • "full": (default) output the full convolution of the inputs.

  • "same": return a centered portion of the "full" output which is the same size as a.

  • "valid": return the portion of the "full" output which do not depend on padding at the array edges.

precision: Specify the precision of the computation. Refer to

jax.lax.Precision for a description of available values.

preferred_element_type: A datatype, indicating to accumulate results to and

return a result with that datatype. Default is None, which means the default accumulation type for the input types.

Returns:

Array containing the convolved result.

See Also:

  • jax.scipy.signal.convolve(): ND convolution

  • jax.numpy.correlate(): 1D correlation

Examples:

A few 1D convolution examples:

>>> x = jnp.array([1, 2, 3, 2, 1])
>>> y = jnp.array([4, 1, 2])

jax.numpy.convolve, by default, returns full convolution using implicit zero-padding at the edges:

>>> jnp.convolve(x, y)
Array([ 4.,  9., 16., 15., 12.,  5.,  2.], dtype=float32)

Specifying mode = 'same' returns a centered convolution the same size as the first input:

>>> jnp.convolve(x, y, mode='same')
Array([ 9., 16., 15., 12.,  5.], dtype=float32)

Specifying mode = 'valid' returns only the portion where the two arrays fully overlap:

>>> jnp.convolve(x, y, mode='valid')
Array([16., 15., 12.], dtype=float32)

For complex-valued inputs:

>>> x1 = jnp.array([3+1j, 2, 4-3j])
>>> y1 = jnp.array([1, 2-3j, 4+5j])
>>> jnp.convolve(x1, y1)
Array([ 3. +1.j, 11. -7.j, 15.+10.j,  7. -8.j, 31. +8.j], dtype=complex64)
Parameters:
Return type:

Array

quchip.declarative.qnp.copy(a, order=None)

Return a copy of the array.

JAX implementation of numpy.copy().

Args:

a: arraylike object to copy order: not implemented in JAX

Returns:

a copy of the input array a.

See Also:

  • jax.numpy.array(): create an array with or without a copy.

  • jax.Array.copy(): same function accessed as an array method.

Examples:

Since JAX arrays are immutable, in most cases explicit array copies are not necessary. One exception is when using a function with donated arguments (see the donate_argnums argument to jax.jit()).

>>> f = jax.jit(lambda x: 2 * x, donate_argnums=0)
>>> x = jnp.arange(4)
>>> y = f(x)
>>> print(y)
[0 2 4 6]

Because we marked x as being donated, the original array is no longer available:

>>> print(x)
Traceback (most recent call last):
RuntimeError: Array has been deleted with shape=int32[4].

In situations like this, an explicit copy will let you keep access to the original buffer:

>>> x = jnp.arange(4)
>>> y = f(x.copy())
>>> print(y)
[0 2 4 6]
>>> print(x)
[0 1 2 3]
Parameters:
Return type:

Array

quchip.declarative.qnp.copysign(x1, x2, /)

Copies the sign of each element in x2 to the corresponding element in x1.

JAX implementation of numpy.copysign.

Args:

x1: Input array x2: The array whose elements will be used to determine the sign, must be

broadcast-compatible with x1

Returns:

An array object containing the potentially changed elements of x1, always promotes to inexact dtype, and has a shape of jnp.broadcast_shapes(x1.shape, x2.shape)

Examples:
>>> x1 = jnp.array([5, 2, 0])
>>> x2 = -1
>>> jnp.copysign(x1, x2)
Array([-5., -2., -0.], dtype=float32)

>>> x1 = jnp.array([6, 8, 0])
>>> x2 = 2
>>> jnp.copysign(x1, x2)
Array([6., 8., 0.], dtype=float32)

>>> x1 = jnp.array([2, -3])
>>> x2 = jnp.array([[1],[-4], [5]])
>>> jnp.copysign(x1, x2)
Array([[ 2.,  3.],
       [-2., -3.],
       [ 2.,  3.]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.corrcoef(x, y=None, rowvar=True)

Compute the Pearson correlation coefficients.

JAX implementation of numpy.corrcoef().

This is a normalized version of the sample covariance computed by jax.numpy.cov(). For a sample covariance \(C_{ij}\), the correlation coefficients are

\[R_{ij} = \frac{C_{ij}}{\sqrt{C_{ii}C_{jj}}}\]

they are constructed such that the values satisfy \(-1 \le R_{ij} \le 1\).

Args:
x: array of shape (M, N) (if rowvar is True), or (N, M)

(if rowvar is False) representing N observations of M variables. x may also be one-dimensional, representing N observations of a single variable.

y: optional set of additional observations, with the same form as m. If

specified, then y is combined with m, i.e. for the default rowvar = True case, m becomes jnp.vstack([m, y]).

rowvar: if True (default) then each row of m represents a variable. If

False, then each column represents a variable.

Returns:

A covariance matrix of shape (M, M).

See also:

  • jax.numpy.cov(): compute the covariance matrix.

Examples:

Consider these observations of two variables that correlate perfectly. The correlation matrix in this case is a 2x2 matrix of ones:

>>> x = jnp.array([[0, 1, 2],
...                [0, 1, 2]])
>>> jnp.corrcoef(x)
Array([[1., 1.],
       [1., 1.]], dtype=float32)

Now consider these observations of two variables that are perfectly anti-correlated. The correlation matrix in this case has -1 in the off-diagonal:

>>> x = jnp.array([[-1,  0,  1],
...                [ 1,  0, -1]])
>>> jnp.corrcoef(x)
Array([[ 1., -1.],
       [-1.,  1.]], dtype=float32)

Equivalently, these sequences can be specified as separate arguments, in which case they are stacked before continuing the computation.

>>> x = jnp.array([-1, 0, 1])
>>> y = jnp.array([1, 0, -1])
>>> jnp.corrcoef(x, y)
Array([[ 1., -1.],
       [-1.,  1.]], dtype=float32)

The entries of the correlation matrix are normalized such that they lie within the range -1 to +1, where +1 indicates perfect correlation and -1 indicates perfect anti-correlation. For example, here is the correlation of 100 points drawn from a 3-dimensional standard normal distribution:

>>> key = jax.random.key(0)
>>> x = jax.random.normal(key, shape=(3, 100))
>>> with jnp.printoptions(precision=2):
...   print(jnp.corrcoef(x))
[[1.   0.03 0.12]
 [0.03 1.   0.01]
 [0.12 0.01 1.  ]]
Parameters:
Return type:

Array

quchip.declarative.qnp.correlate(a, v, mode='valid', *, precision=None, preferred_element_type=None)

Correlation of two one dimensional arrays.

JAX implementation of numpy.correlate().

Correlation of one dimensional arrays is defined as:

\[c_k = \sum_j a_{k + j} \overline{v_j}\]

where \(\overline{v_j}\) is the complex conjugate of \(v_j\).

Args:

a: left-hand input to the correlation. Must have a.ndim == 1. v: right-hand input to the correlation. Must have v.ndim == 1. mode: controls the size of the output. Available operations are:

  • "full": output the full correlation of the inputs.

  • "same": return a centered portion of the "full" output which is the same size as a.

  • "valid": (default) return the portion of the "full" output which do not depend on padding at the array edges.

precision: Specify the precision of the computation. Refer to

jax.lax.Precision for a description of available values.

preferred_element_type: A datatype, indicating to accumulate results to and

return a result with that datatype. Default is None, which means the default accumulation type for the input types.

Returns:

Array containing the cross-correlation result.

See Also:

  • jax.scipy.signal.correlate(): ND correlation

  • jax.numpy.convolve(): 1D convolution

Examples:
>>> x = jnp.array([1, 2, 3, 2, 1])
>>> y = jnp.array([4, 5, 6])

Since default mode = 'valid', jax.numpy.correlate returns only the portion of correlation where the two arrays fully overlap:

>>> jnp.correlate(x, y)
Array([32., 35., 28.], dtype=float32)

Specifying mode = 'full' returns full correlation using implicit zero-padding at the edges.

>>> jnp.correlate(x, y, mode='full')
Array([ 6., 17., 32., 35., 28., 13.,  4.], dtype=float32)

Specifying mode = 'same' returns a centered correlation the same size as the first input:

>>> jnp.correlate(x, y, mode='same')
Array([17., 32., 35., 28., 13.], dtype=float32)

If both the inputs arrays are real-valued and symmetric then the result will also be symmetric and will be equal to the result of jax.numpy.convolve.

>>> x1 = jnp.array([1, 2, 3, 2, 1])
>>> y1 = jnp.array([4, 5, 4])
>>> jnp.correlate(x1, y1, mode='full')
Array([ 4., 13., 26., 31., 26., 13.,  4.], dtype=float32)
>>> jnp.convolve(x1, y1, mode='full')
Array([ 4., 13., 26., 31., 26., 13.,  4.], dtype=float32)

For complex-valued inputs:

>>> x2 = jnp.array([3+1j, 2, 2-3j])
>>> y2 = jnp.array([4, 2-5j, 1])
>>> jnp.correlate(x2, y2, mode='full')
Array([ 3. +1.j,  3.+17.j, 18.+11.j, 27. +4.j,  8.-12.j], dtype=complex64)
Parameters:
Return type:

Array

quchip.declarative.qnp.cos(x, /)

Compute a trigonometric cosine of each element of input.

JAX implementation of numpy.cos.

Args:

x: scalar or array. Angle in radians.

Returns:

An array containing the cosine of each element in x, promotes to inexact dtype.

See also:

  • jax.numpy.sin(): Computes a trigonometric sine of each element of input.

  • jax.numpy.tan(): Computes a trigonometric tangent of each element of input.

  • jax.numpy.arccos() and jax.numpy.acos(): Computes the inverse of trigonometric cosine of each element of input.

Examples:
>>> pi = jnp.pi
>>> x = jnp.array([pi/4, pi/2, 3*pi/4, 5*pi/6])
>>> with jnp.printoptions(precision=3, suppress=True):
...   print(jnp.cos(x))
[ 0.707 -0.    -0.707 -0.866]
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.cosh(x, /)

Calculate element-wise hyperbolic cosine of input.

JAX implementation of numpy.cosh.

The hyperbolic cosine is defined by:

\[cosh(x) = \frac{e^x + e^{-x}}{2}\]
Args:

x: input array or scalar.

Returns:

An array containing the hyperbolic cosine of each element of x, promoting to inexact dtype.

Note:

jnp.cosh is equivalent to computing jnp.cos(1j * x).

See also:

  • jax.numpy.sinh(): Computes the element-wise hyperbolic sine of the input.

  • jax.numpy.tanh(): Computes the element-wise hyperbolic tangent of the input.

  • jax.numpy.arccosh(): Computes the element-wise inverse of hyperbolic cosine of the input.

Examples:
>>> x = jnp.array([[3, -1, 0],
...                [4, 7, -5]])
>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.cosh(x)
Array([[ 10.068,   1.543,   1.   ],
       [ 27.308, 548.317,  74.21 ]], dtype=float32)
>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.cos(1j * x)
Array([[ 10.068+0.j,   1.543+0.j,   1.   +0.j],
       [ 27.308+0.j, 548.317+0.j,  74.21 +0.j]],      dtype=complex64, weak_type=True)

For complex-valued input:

>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.cosh(5+1j)
Array(40.096+62.44j, dtype=complex64, weak_type=True)
>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.cos(1j * (5+1j))
Array(40.096+62.44j, dtype=complex64, weak_type=True)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.count_nonzero(a, axis=None, keepdims=False)

Return the number of nonzero elements along a given axis.

JAX implementation of numpy.count_nonzero().

Args:

a: input array. axis: optional, int or sequence of ints, default=None. Axis along which the

number of nonzeros are counted. If None, counts within the flattened array.

keepdims: bool, default=False. If true, reduced axes are left in the result

with size 1.

Returns:

An array with number of nonzeros elements along specified axis of the input.

Examples:

By default, jnp.count_nonzero counts the nonzero values along all axes.

>>> x = jnp.array([[1, 0, 0, 0],
...                [0, 0, 1, 0],
...                [1, 1, 1, 0]])
>>> jnp.count_nonzero(x)
Array(5, dtype=int32)

If axis=1, counts along axis 1.

>>> jnp.count_nonzero(x, axis=1)
Array([1, 1, 3], dtype=int32)

To preserve the dimensions of input, you can set keepdims=True.

>>> jnp.count_nonzero(x, axis=1, keepdims=True)
Array([[1],
       [1],
       [3]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.cov(m, y=None, rowvar=True, bias=False, ddof=None, fweights=None, aweights=None)

Estimate the weighted sample covariance.

JAX implementation of numpy.cov().

The covariance \(C_{ij}\) between variable i and variable j is defined as

\[cov[X_i, X_j] = E[(X_i - E[X_i])(X_j - E[X_j])]\]

Given an array of N observations of the variables \(X_i\) and \(X_j\), this can be estimated via the sample covariance:

\[C_{ij} = \frac{1}{N - 1} \sum_{n=1}^N (X_{in} - \overline{X_i})(X_{jn} - \overline{X_j})\]

Where \(\overline{X_i} = \frac{1}{N} \sum_{k=1}^N X_{ik}\) is the mean of the observations.

Args:
m: array of shape (M, N) (if rowvar is True), or (N, M)

(if rowvar is False) representing N observations of M variables. m may also be one-dimensional, representing N observations of a single variable.

y: optional set of additional observations, with the same form as m. If

specified, then y is combined with m, i.e. for the default rowvar = True case, m becomes jnp.vstack([m, y]).

rowvar: if True (default) then each row of m represents a variable. If

False, then each column represents a variable.

bias: if False (default) then normalize the covariance by N - 1. If True,

then normalize the covariance by N

ddof: specify the degrees of freedom. Defaults to 1 if bias is False,

or to 0 if bias is True.

fweights: optional array of integer frequency weights of shape (N,). This

is an absolute weight specifying the number of times each observation is included in the computation.

aweights: optional array of observation weights of shape (N,). This is

a relative weight specifying the “importance” of each observation. In the ddof=0 case, it is equivalent to assigning probabilities to each observation.

Returns:

A covariance matrix of shape (M, M), or a scalar with shape () if M = 1.

See also:

  • jax.numpy.corrcoef(): compute the correlation coefficient, a normalized version of the covariance matrix.

Examples:

Consider these observations of two variables that correlate perfectly. The covariance matrix in this case is a 2x2 matrix of ones:

>>> x = jnp.array([[0, 1, 2],
...                [0, 1, 2]])
>>> jnp.cov(x)
Array([[1., 1.],
       [1., 1.]], dtype=float32)

Now consider these observations of two variables that are perfectly anti-correlated. The covariance matrix in this case has -1 in the off-diagonal:

>>> x = jnp.array([[-1,  0,  1],
...                [ 1,  0, -1]])
>>> jnp.cov(x)
Array([[ 1., -1.],
       [-1.,  1.]], dtype=float32)

Equivalently, these sequences can be specified as separate arguments, in which case they are stacked before continuing the computation.

>>> x = jnp.array([-1, 0, 1])
>>> y = jnp.array([1, 0, -1])
>>> jnp.cov(x, y)
Array([[ 1., -1.],
       [-1.,  1.]], dtype=float32)

In general, the entries of the covariance matrix may be any positive or negative real value. For example, here is the covariance of 100 points drawn from a 3-dimensional standard normal distribution:

>>> key = jax.random.key(0)
>>> x = jax.random.normal(key, shape=(3, 100))
>>> with jnp.printoptions(precision=2):
...   print(jnp.cov(x))
[[0.9  0.03 0.1 ]
 [0.03 1.   0.01]
 [0.1  0.01 0.85]]
Parameters:
Return type:

Array

quchip.declarative.qnp.cross(a, b, axisa=-1, axisb=-1, axisc=-1, axis=None)

Compute the (batched) cross product of two arrays.

JAX implementation of numpy.cross().

This computes the 2-dimensional or 3-dimensional cross product,

\[c = a \times b\]

In 3 dimensions, c is a length-3 array. In 2 dimensions, c is a scalar.

Args:
a: N-dimensional array. a.shape[axisa] indicates the dimension of

the cross product, and must be 2 or 3.

b: N-dimensional array. Must have b.shape[axisb] == a.shape[axisb],

and other dimensions of a and b must be broadcast compatible.

axisa: specicy the axis of a along which to compute the cross product. axisb: specicy the axis of b along which to compute the cross product. axisc: specicy the axis of c along which the cross product result

will be stored.

axis: if specified, this overrides axisa, axisb, and axisc

with a single value.

Returns:

The array c containing the (batched) cross product of a and b along the specified axes.

See also:

  • jax.numpy.linalg.cross(): an array API compatible function for computing cross products over 3-vectors.

Examples:

A 2-dimensional cross product returns a scalar:

>>> a = jnp.array([1, 2])
>>> b = jnp.array([3, 4])
>>> jnp.cross(a, b)
Array(-2, dtype=int32)

A 3-dimensional cross product returns a length-3 vector:

>>> a = jnp.array([1, 2, 3])
>>> b = jnp.array([4, 5, 6])
>>> jnp.cross(a, b)
Array([-3,  6, -3], dtype=int32)

With multi-dimensional inputs, the cross-product is computed along the last axis by default. Here’s a batched 3-dimensional cross product, operating on the rows of the inputs:

>>> a = jnp.array([[1, 2, 3],
...                [3, 4, 3]])
>>> b = jnp.array([[2, 3, 2],
...                [4, 5, 6]])
>>> jnp.cross(a, b)
Array([[-5,  4, -1],
       [ 9, -6, -1]], dtype=int32)

Specifying axis=0 makes this a batched 2-dimensional cross product, operating on the columns of the inputs:

>>> jnp.cross(a, b, axis=0)
Array([-2, -2, 12], dtype=int32)

Equivalently, we can independently specify the axis of the inputs a and b and the output c:

>>> jnp.cross(a, b, axisa=0, axisb=0, axisc=0)
Array([-2, -2, 12], dtype=int32)
Parameters:
quchip.declarative.qnp.csingle

alias of complex64

quchip.declarative.qnp.cumprod(a, axis=None, dtype=None, out=None)

Cumulative product of elements along an axis.

JAX implementation of numpy.cumprod().

Args:

a: N-dimensional array to be accumulated. axis: integer axis along which to accumulate. If None (default), then

array will be flattened and accumulated along the flattened axis.

dtype: optionally specify the dtype of the output. If not specified,

then the output dtype will match the input dtype.

out: unused by JAX

Returns:

An array containing the accumulated product along the given axis.

See also:

  • jax.numpy.multiply.accumulate(): cumulative product via ufunc methods.

  • jax.numpy.nancumprod(): cumulative product ignoring NaN values.

  • jax.numpy.prod(): product along axis

Examples:
>>> x = jnp.array([[1, 2, 3],
...                [4, 5, 6]])
>>> jnp.cumprod(x)  # flattened cumulative product
Array([  1,   2,   6,  24, 120, 720], dtype=int32)
>>> jnp.cumprod(x, axis=1)  # cumulative product along axis 1
Array([[  1,   2,   6],
       [  4,  20, 120]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.cumsum(a, axis=None, dtype=None, out=None)

Cumulative sum of elements along an axis.

JAX implementation of numpy.cumsum().

Args:

a: N-dimensional array to be accumulated. axis: integer axis along which to accumulate. If None (default), then

array will be flattened and accumulated along the flattened axis.

dtype: optionally specify the dtype of the output. If not specified,

then the output dtype will match the input dtype.

out: unused by JAX

Returns:

An array containing the accumulated sum along the given axis.

See also:

  • jax.numpy.cumulative_sum(): cumulative sum via the array API standard.

  • jax.numpy.add.accumulate(): cumulative sum via ufunc methods.

  • jax.numpy.nancumsum(): cumulative sum ignoring NaN values.

  • jax.numpy.sum(): sum along axis

Examples:
>>> x = jnp.array([[1, 2, 3],
...                [4, 5, 6]])
>>> jnp.cumsum(x)  # flattened cumulative sum
Array([ 1,  3,  6, 10, 15, 21], dtype=int32)
>>> jnp.cumsum(x, axis=1)  # cumulative sum along axis 1
Array([[ 1,  3,  6],
       [ 4,  9, 15]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.cumulative_prod(x, /, *, axis=None, dtype=None, include_initial=False)

Cumulative product along the axis of an array.

JAX implementation of numpy.cumulative_prod().

Args:

x: N-dimensional array axis: integer axis along which to accumulate. If x is one-dimensional,

this argument is optional and defaults to zero.

dtype: optional dtype of the output. include_initial: if True, then include the initial value in the cumulative

product. Default is False.

Returns:

An array containing the accumulated values.

See Also:

  • jax.numpy.cumprod(): alternative API for cumulative product.

  • jax.numpy.nancumprod(): cumulative product while ignoring NaN values.

  • jax.numpy.multiply.accumulate(): cumulative product via the ufunc API.

Examples:
>>> x = jnp.array([[1, 2, 3],
...                [4, 5, 6]])
>>> jnp.cumulative_prod(x, axis=1)
Array([[  1,   2,   6],
       [  4,  20, 120]], dtype=int32)
>>> jnp.cumulative_prod(x, axis=1, include_initial=True)
Array([[  1,   1,   2,   6],
       [  1,   4,  20, 120]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.cumulative_sum(x, /, *, axis=None, dtype=None, include_initial=False)

Cumulative sum along the axis of an array.

JAX implementation of numpy.cumulative_sum().

Args:

x: N-dimensional array axis: integer axis along which to accumulate. If x is one-dimensional,

this argument is optional and defaults to zero.

dtype: optional dtype of the output. include_initial: if True, then include the initial value in the cumulative

sum. Default is False.

Returns:

An array containing the accumulated values.

See Also:

  • jax.numpy.cumsum(): alternative API for cumulative sum.

  • jax.numpy.nancumsum(): cumulative sum while ignoring NaN values.

  • jax.numpy.add.accumulate(): cumulative sum via the ufunc API.

Examples:
>>> x = jnp.array([[1, 2, 3],
...                [4, 5, 6]])
>>> jnp.cumulative_sum(x, axis=1)
Array([[ 1,  3,  6],
       [ 4,  9, 15]], dtype=int32)
>>> jnp.cumulative_sum(x, axis=1, include_initial=True)
Array([[ 0,  1,  3,  6],
       [ 0,  4,  9, 15]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.deg2rad(x, /)

Convert angles from degrees to radians.

JAX implementation of numpy.deg2rad.

The angle in degrees is converted to radians by:

\[deg2rad(x) = x * \frac{pi}{180}\]
Args:

x: scalar or array. Specifies the angle in degrees.

Returns:

An array containing the angles in radians.

See also:

  • jax.numpy.rad2deg() and jax.numpy.degrees(): Converts the angles from radians to degrees.

  • jax.numpy.radians(): Alias of deg2rad.

Examples:
>>> x = jnp.array([60, 90, 120, 180])
>>> jnp.deg2rad(x)
Array([1.0471976, 1.5707964, 2.0943952, 3.1415927], dtype=float32)
>>> x * jnp.pi / 180
Array([1.0471976, 1.5707964, 2.0943952, 3.1415927],      dtype=float32, weak_type=True)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.degrees(x, /)

Alias of jax.numpy.rad2deg()

Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.delete(arr, obj, axis=None, *, assume_unique_indices=False)

Delete entry or entries from an array.

JAX implementation of numpy.delete().

Args:

arr: array from which entries will be deleted. obj: index, indices, or slice to be deleted. axis: axis along which entries will be deleted. assume_unique_indices: In case of array-like integer (not boolean) indices,

assume the indices are unique, and perform the deletion in a way that is compatible with JIT and other JAX transformations.

Returns:

Copy of arr with specified indices deleted.

Note:

delete() usually requires the index specification to be static. If the index is an integer array that is guaranteed to contain unique entries, you may specify assume_unique_indices=True to perform the operation in a manner that does not require static indices.

See also:

  • jax.numpy.insert(): insert entries into an array.

Examples:

Delete entries from a 1D array:

>>> a = jnp.array([4, 5, 6, 7, 8, 9])
>>> jnp.delete(a, 2)
Array([4, 5, 7, 8, 9], dtype=int32)
>>> jnp.delete(a, slice(1, 4))  # delete a[1:4]
Array([4, 8, 9], dtype=int32)
>>> jnp.delete(a, slice(None, None, 2))  # delete a[::2]
Array([5, 7, 9], dtype=int32)

Delete entries from a 2D array along a specified axis:

>>> a2 = jnp.array([[4, 5, 6],
...                 [7, 8, 9]])
>>> jnp.delete(a2, 1, axis=1)
Array([[4, 6],
       [7, 9]], dtype=int32)

Delete multiple entries via a sequence of indices:

>>> indices = jnp.array([0, 1, 3])
>>> jnp.delete(a, indices)
Array([6, 8, 9], dtype=int32)

This will fail under jit() and other transformations, because the output shape cannot be known with the possibility of duplicate indices:

>>> jax.jit(jnp.delete)(a, indices)
Traceback (most recent call last):
  ...
ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: traced array with shape int32[3].

If you can ensure that the indices are unique, pass assume_unique_indices to allow this to be executed under JIT:

>>> jit_delete = jax.jit(jnp.delete, static_argnames=['assume_unique_indices'])
>>> jit_delete(a, indices, assume_unique_indices=True)
Array([6, 8, 9], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.diag(v, k=0)

Returns the specified diagonal or constructs a diagonal array.

JAX implementation of numpy.diag().

The JAX version always returns a copy of the input, although if this is used within a JIT compilation, the compiler may avoid the copy.

Args:
v: Input array. Can be a 1-D array to create a diagonal matrix or a

2-D array to extract a diagonal.

k: optional, default=0. Diagonal offset. Positive values place the diagonal

above the main diagonal, negative values place it below the main diagonal.

Returns:

If v is a 2-D array, a 1-D array containing the diagonal elements. If v is a 1-D array, a 2-D array with the input elements placed along the specified diagonal.

See also:

  • jax.numpy.diagflat()

  • jax.numpy.diagonal()

Examples:

Creating a diagonal matrix from a 1-D array:

>>> jnp.diag(jnp.array([1, 2, 3]))
Array([[1, 0, 0],
       [0, 2, 0],
       [0, 0, 3]], dtype=int32)

Specifying a diagonal offset:

>>> jnp.diag(jnp.array([1, 2, 3]), k=1)
Array([[0, 1, 0, 0],
       [0, 0, 2, 0],
       [0, 0, 0, 3],
       [0, 0, 0, 0]], dtype=int32)

Extracting a diagonal from a 2-D array:

>>> x = jnp.array([[1, 2, 3],
...                [4, 5, 6],
...                [7, 8, 9]])
>>> jnp.diag(x)
Array([1, 5, 9], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.diag_indices(n, ndim=2)

Return indices for accessing the main diagonal of a multidimensional array.

JAX implementation of numpy.diag_indices().

Args:

n: int. The size of each dimension of the square array. ndim: optional, int, default=2. The number of dimensions of the array.

Returns:

A tuple of arrays, each of length n, containing the indices to access the main diagonal.

See also:

  • jax.numpy.diag_indices_from()

  • jax.numpy.diagonal()

Examples:
>>> jnp.diag_indices(3)
(Array([0, 1, 2], dtype=int32), Array([0, 1, 2], dtype=int32))
>>> jnp.diag_indices(4, ndim=3)
(Array([0, 1, 2, 3], dtype=int32),
Array([0, 1, 2, 3], dtype=int32),
Array([0, 1, 2, 3], dtype=int32))
Parameters:
Return type:

tuple[Array, …]

quchip.declarative.qnp.diag_indices_from(arr)

Return indices for accessing the main diagonal of a given array.

JAX implementation of numpy.diag_indices_from().

Args:
arr: Input array. Must be at least 2-dimensional and have equal length along

all dimensions.

Returns:

A tuple of arrays containing the indices to access the main diagonal of the input array.

See also:

  • jax.numpy.diag_indices()

  • jax.numpy.diagonal()

Examples:
>>> arr = jnp.array([[1, 2, 3],
...                  [4, 5, 6],
...                  [7, 8, 9]])
>>> jnp.diag_indices_from(arr)
(Array([0, 1, 2], dtype=int32), Array([0, 1, 2], dtype=int32))
>>> arr = jnp.array([[[1, 2], [3, 4]],
...                  [[5, 6], [7, 8]]])
>>> jnp.diag_indices_from(arr)
(Array([0, 1], dtype=int32),
Array([0, 1], dtype=int32),
Array([0, 1], dtype=int32))
Parameters:

arr (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

tuple[Array, …]

quchip.declarative.qnp.diagflat(v, k=0)

Return a 2-D array with the flattened input array laid out on the diagonal.

JAX implementation of numpy.diagflat().

This differs from np.diagflat for some scalar values of v. JAX always returns a two-dimensional array, whereas NumPy may return a scalar depending on the type of v.

Args:

v: Input array. Can be N-dimensional but is flattened to 1D. k: optional, default=0. Diagonal offset. Positive values place the diagonal

above the main diagonal, negative values place it below the main diagonal.

Returns:

A 2D array with the input elements placed along the diagonal with the specified offset (k). The remaining entries are filled with zeros.

See also:

  • jax.numpy.diag()

  • jax.numpy.diagonal()

Examples:
>>> jnp.diagflat(jnp.array([1, 2, 3]))
Array([[1, 0, 0],
       [0, 2, 0],
       [0, 0, 3]], dtype=int32)
>>> jnp.diagflat(jnp.array([1, 2, 3]), k=1)
Array([[0, 1, 0, 0],
       [0, 0, 2, 0],
       [0, 0, 0, 3],
       [0, 0, 0, 0]], dtype=int32)
>>> a = jnp.array([[1, 2],
...                [3, 4]])
>>> jnp.diagflat(a)
Array([[1, 0, 0, 0],
       [0, 2, 0, 0],
       [0, 0, 3, 0],
       [0, 0, 0, 4]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.diagonal(a, offset=0, axis1=0, axis2=1)

Returns the specified diagonal of an array.

JAX implementation of numpy.diagonal().

The JAX version always returns a copy of the input, although if this is used within a JIT compilation, the compiler may avoid the copy.

Args:

a: Input array. Must be at least 2-dimensional. offset: optional, default=0. Diagonal offset from the main diagonal.

Must be a static integer value. Can be positive or negative.

axis1: optional, default=0. The first axis along which to take the diagonal. axis2: optional, default=1. The second axis along which to take the diagonal.

Returns:

A 1D array for 2D input, and in general a N-1 dimensional array for N-dimensional input.

See also:

  • jax.numpy.diag()

  • jax.numpy.diagflat()

Examples:
>>> x = jnp.array([[1, 2, 3],
...                [4, 5, 6],
...                [7, 8, 9]])
>>> jnp.diagonal(x)
Array([1, 5, 9], dtype=int32)
>>> jnp.diagonal(x, offset=1)
Array([2, 6], dtype=int32)
>>> jnp.diagonal(x, offset=-1)
Array([4, 8], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.diff(a, n=1, axis=-1, prepend=None, append=None)

Calculate n-th order difference between array elements along a given axis.

JAX implementation of numpy.diff().

The first order difference is computed by a[i+1] - a[i], and the n-th order difference is computed n times recursively.

Args:

a: input array. Must have a.ndim >= 1. n: int, optional, default=1. Order of the difference. Specifies the number

of times the difference is computed. If n=0, no difference is computed and input is returned as is.

axis: int, optional, default=-1. Specifies the axis along which the difference

is computed. The difference is computed along axis -1 by default.

prepend: scalar or array, optional, default=None. Specifies the values to be

prepended along axis before computing the difference.

append: scalar or array, optional, default=None. Specifies the values to be

appended along axis before computing the difference.

Returns:

An array containing the n-th order difference between the elements of a.

See also:

  • jax.numpy.ediff1d(): Computes the differences between consecutive elements of an array.

  • jax.numpy.cumsum(): Computes the cumulative sum of the elements of the array along a given axis.

  • jax.numpy.gradient(): Computes the gradient of an N-dimensional array.

Examples:

jnp.diff computes the first order difference along axis, by default.

>>> a = jnp.array([[1, 5, 2, 9],
...                [3, 8, 7, 4]])
>>> jnp.diff(a)
Array([[ 4, -3,  7],
       [ 5, -1, -3]], dtype=int32)

When n = 2, second order difference is computed along axis.

>>> jnp.diff(a, n=2)
Array([[-7, 10],
       [-6, -2]], dtype=int32)

When prepend = 2, it is prepended to a along axis before computing the difference.

>>> jnp.diff(a, prepend=2)
Array([[-1,  4, -3,  7],
       [ 1,  5, -1, -3]], dtype=int32)

When append = jnp.array([[3],[1]]), it is appended to a along axis before computing the difference.

>>> jnp.diff(a, append=jnp.array([[3],[1]]))
Array([[ 4, -3,  7, -6],
       [ 5, -1, -3, -3]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.digitize(x, bins, right=False, *, method=None)

Convert an array to bin indices.

JAX implementation of numpy.digitize().

Args:

x: array of values to digitize. bins: 1D array of bin edges. Must be monotonically increasing or decreasing. right: if true, the intervals include the right bin edges. If false (default)

the intervals include the left bin edges.

method: optional method argument to be passed to searchsorted().

See that function for available options.

Returns:

An integer array of the same shape as x indicating the bin number that the values are in.

See also:

  • jax.numpy.searchsorted(): find insertion indices for values in a sorted array.

  • jax.numpy.histogram(): compute frequency of array values within specified bins.

Examples:
>>> x = jnp.array([1.0, 2.0, 2.5, 1.5, 3.0, 3.5])
>>> bins = jnp.array([1, 2, 3])
>>> jnp.digitize(x, bins)
Array([1, 2, 2, 1, 3, 3], dtype=int32)
>>> jnp.digitize(x, bins, right=True)
Array([0, 1, 2, 1, 2, 3], dtype=int32)

digitize supports reverse-ordered bins as well:

>>> bins = jnp.array([3, 2, 1])
>>> jnp.digitize(x, bins)
Array([2, 1, 1, 2, 0, 0], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.divide(x1, x2, /)

Alias of jax.numpy.true_divide().

Parameters:
Return type:

Array

quchip.declarative.qnp.divmod(x1, x2, /)

Calculates the integer quotient and remainder of x1 by x2 element-wise

JAX implementation of numpy.divmod.

Args:

x1: Input array, the dividend x2: Input array, the divisor

Returns:

A tuple of arrays (x1 // x2, x1 % x2).

See Also:

  • jax.numpy.floor_divide(): floor division function

  • jax.numpy.remainder(): remainder function

Examples:
>>> x1 = jnp.array([10, 20, 30])
>>> x2 = jnp.array([3, 4, 7])
>>> jnp.divmod(x1, x2)
(Array([3, 5, 4], dtype=int32), Array([1, 0, 2], dtype=int32))

>>> x1 = jnp.array([-5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5])
>>> x2 = 3
>>> jnp.divmod(x1, x2)
(Array([-2, -2, -1, -1, -1,  0,  0,  0,  1,  1,  1], dtype=int32),
 Array([1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2], dtype=int32))

>>> x1 = jnp.array([6, 6, 6], dtype=jnp.int32)
>>> x2 = jnp.array([1.9, 2.5, 3.1], dtype=jnp.float32)
>>> jnp.divmod(x1, x2)
(Array([3., 2., 1.], dtype=float32),
 Array([0.30000007, 1.        , 2.9       ], dtype=float32))
Parameters:
Return type:

tuple[Array, Array]

quchip.declarative.qnp.dot(a, b, *, precision=None, preferred_element_type=None, out_sharding=None)

Compute the dot product of two arrays.

JAX implementation of numpy.dot().

This differs from jax.numpy.matmul() in two respects:

  • if either a or b is a scalar, the result of dot is equivalent to jax.numpy.multiply(), while the result of matmul is an error.

  • if a and b have more than 2 dimensions, the batch indices are stacked rather than broadcast.

Args:

a: first input array, of shape (..., N). b: second input array. Must have shape (N,) or (..., N, M).

In the multi-dimensional case, leading dimensions must be broadcast-compatible with the leading dimensions of a.

precision: either None (default), which means the default precision for

the backend, a Precision enum value (Precision.DEFAULT, Precision.HIGH or Precision.HIGHEST) or a tuple of two such values indicating precision of a and b.

preferred_element_type: either None (default), which means the default

accumulation type for the input types, or a datatype, indicating to accumulate results to and return a result with that datatype.

Returns:

array containing the dot product of the inputs, with batch dimensions of a and b stacked rather than broadcast.

See also:

  • jax.numpy.matmul(): broadcasted batched matmul.

  • jax.lax.dot_general(): general batched matrix multiplication.

Examples:

For scalar inputs, dot computes the element-wise product:

>>> x = jnp.array([1, 2, 3])
>>> jnp.dot(x, 2)
Array([2, 4, 6], dtype=int32)

For vector or matrix inputs, dot computes the vector or matrix product:

>>> M = jnp.array([[2, 3, 4],
...                [5, 6, 7],
...                [8, 9, 0]])
>>> jnp.dot(M, x)
Array([20, 38, 26], dtype=int32)
>>> jnp.dot(M, M)
Array([[ 51,  60,  29],
       [ 96, 114,  62],
       [ 61,  78,  95]], dtype=int32)

For higher-dimensional matrix products, batch dimensions are stacked, whereas in matmul() they are broadcast. For example:

>>> a = jnp.zeros((3, 2, 4))
>>> b = jnp.zeros((3, 4, 1))
>>> jnp.dot(a, b).shape
(3, 2, 3, 1)
>>> jnp.matmul(a, b).shape
(3, 2, 1)
Parameters:
Return type:

Array

quchip.declarative.qnp.double

alias of float64

quchip.declarative.qnp.dsplit(ary, indices_or_sections)

Split an array into sub-arrays depth-wise.

JAX implementation of numpy.dsplit().

Refer to the documentation of jax.numpy.split() for details. dsplit is equivalent to split with axis=2.

Examples:

>>> x = jnp.arange(12).reshape(3, 1, 4)
>>> print(x)
[[[ 0  1  2  3]]

 [[ 4  5  6  7]]

 [[ 8  9 10 11]]]
>>> x1, x2 = jnp.dsplit(x, 2)
>>> print(x1)
[[[0 1]]

 [[4 5]]

 [[8 9]]]
>>> print(x2)
[[[ 2  3]]

 [[ 6  7]]

 [[10 11]]]
See also:

  • jax.numpy.split(): split an array along any axis.

  • jax.numpy.vsplit(): split vertically, i.e. along axis=0

  • jax.numpy.hsplit(): split horizontally, i.e. along axis=1

  • jax.numpy.array_split(): like split, but allows indices_or_sections to be an integer that does not evenly divide the size of the array.

Parameters:
Return type:

list[Array]

quchip.declarative.qnp.dstack(tup, dtype=None)

Stack arrays depth-wise.

JAX implementation of numpy.dstack().

For arrays of three or more dimensions, this is equivalent to jax.numpy.concatenate() with axis=2.

Args:
tup: a sequence of arrays to stack; each must have the same shape along all

but the third axis. Input arrays will be promoted to at least rank 3. If a single array is given it will be treated equivalently to tup = unstack(tup), but the implementation will avoid explicit unstacking.

dtype: optional dtype of the resulting array. If not specified, the dtype

will be determined via type promotion rules described in type-promotion.

Returns:

the stacked result.

See also:

  • jax.numpy.stack(): stack along arbitrary axes

  • jax.numpy.concatenate(): concatenation along existing axes.

  • jax.numpy.vstack(): stack vertically, i.e. along axis 0.

  • jax.numpy.hstack(): stack horizontally, i.e. along axis 1.

Examples:

Scalar values:

>>> jnp.dstack([1, 2, 3])
Array([[[1, 2, 3]]], dtype=int32, weak_type=True)

1D arrays:

>>> x = jnp.arange(3)
>>> y = jnp.ones(3)
>>> jnp.dstack([x, y])
Array([[[0., 1.],
        [1., 1.],
        [2., 1.]]], dtype=float32)

2D arrays:

>>> x = x.reshape(1, 3)
>>> y = y.reshape(1, 3)
>>> jnp.dstack([x, y])
Array([[[0., 1.],
        [1., 1.],
        [2., 1.]]], dtype=float32)
Parameters:
Return type:

Array

class quchip.declarative.qnp.dtype(dtype, align=False, copy=False[, metadata])

Bases: object

Create a data type object.

A numpy array is homogeneous, and contains elements described by a dtype object. A dtype object can be constructed from different combinations of fundamental numeric types.

Parameters:

  • dtype – Object to be converted to a data type object.

  • align (bool, optional) – Add padding to the fields to match what a C compiler would output for a similar C-struct. Can be True only if obj is a dictionary or a comma-separated string. If a struct dtype is being created, this also sets a sticky alignment flag isalignedstruct.

  • copy (bool, optional) – Make a new copy of the data-type object. If False, the result may just be a reference to a built-in data-type object.

  • metadata (dict, optional) – An optional dictionary with dtype metadata.

See also

result_type

Examples

Using array-scalar type:

>>> import numpy as np
>>> np.dtype(np.int16)
dtype('int16')

Structured type, one field name ‘f1’, containing int16:

>>> np.dtype([('f1', np.int16)])
dtype([('f1', '<i2')])

Structured type, one field named ‘f1’, in itself containing a structured type with one field:

>>> np.dtype([('f1', [('f1', np.int16)])])
dtype([('f1', [('f1', '<i2')])])

Structured type, two fields: the first field contains an unsigned int, the second an int32:

>>> np.dtype([('f1', np.uint64), ('f2', np.int32)])
dtype([('f1', '<u8'), ('f2', '<i4')])

Using array-protocol type strings:

>>> np.dtype([('a','f8'),('b','S10')])
dtype([('a', '<f8'), ('b', 'S10')])

Using comma-separated field formats. The shape is (2,3):

>>> np.dtype("i4, (2,3)f8")
dtype([('f0', '<i4'), ('f1', '<f8', (2, 3))])

Using tuples. int is a fixed type, 3 the field’s shape. void is a flexible type, here of size 10:

>>> np.dtype([('hello',(np.int64,3)),('world',np.void,10)])
dtype([('hello', '<i8', (3,)), ('world', 'V10')])

Subdivide int16 into 2 int8’s, called x and y. 0 and 1 are the offsets in bytes:

>>> np.dtype((np.int16, {'x':(np.int8,0), 'y':(np.int8,1)}))
dtype((numpy.int16, [('x', 'i1'), ('y', 'i1')]))

Using dictionaries. Two fields named ‘gender’ and ‘age’:

>>> np.dtype({'names':['gender','age'], 'formats':['S1',np.uint8]})
dtype([('gender', 'S1'), ('age', 'u1')])

Offsets in bytes, here 0 and 25:

>>> np.dtype({'surname':('S25',0),'age':(np.uint8,25)})
dtype([('surname', 'S25'), ('age', 'u1')])
alignment

The required alignment (bytes) of this data-type according to the compiler.

More information is available in the C-API section of the manual.

Examples

>>> import numpy as np
>>> x = np.dtype('i4')
>>> x.alignment
4

>>> x = np.dtype(float)
>>> x.alignment
8
base

Returns dtype for the base element of the subarrays, regardless of their dimension or shape.

See also

dtype.subdtype

Examples

>>> import numpy as np
>>> x = numpy.dtype('8f')
>>> x.base
dtype('float32')

>>> x =  numpy.dtype('i2')
>>> x.base
dtype('int16')
byteorder

A character indicating the byte-order of this data-type object.

One of:

‘=’

native

‘<’

little-endian

‘>’

big-endian

‘|’

not applicable

All built-in data-type objects have byteorder either ‘=’ or ‘|’.

Examples

>>> import numpy as np
>>> dt = np.dtype('i2')
>>> dt.byteorder
'='
>>> # endian is not relevant for 8 bit numbers
>>> np.dtype('i1').byteorder
'|'
>>> # or ASCII strings
>>> np.dtype('S2').byteorder
'|'
>>> # Even if specific code is given, and it is native
>>> # '=' is the byteorder
>>> import sys
>>> sys_is_le = sys.byteorder == 'little'
>>> native_code = '<' if sys_is_le else '>'
>>> swapped_code = '>' if sys_is_le else '<'
>>> dt = np.dtype(native_code + 'i2')
>>> dt.byteorder
'='
>>> # Swapped code shows up as itself
>>> dt = np.dtype(swapped_code + 'i2')
>>> dt.byteorder == swapped_code
True
char

A unique character code for each of the 21 different built-in types.

Examples

>>> import numpy as np
>>> x = np.dtype(float)
>>> x.char
'd'
descr

__array_interface__ description of the data-type.

The format is that required by the ‘descr’ key in the __array_interface__ attribute.

Warning: This attribute exists specifically for __array_interface__, and passing it directly to numpy.dtype will not accurately reconstruct some dtypes (e.g., scalar and subarray dtypes).

Examples

>>> import numpy as np
>>> x = np.dtype(float)
>>> x.descr
[('', '<f8')]

>>> dt = np.dtype([('name', np.str_, 16), ('grades', np.float64, (2,))])
>>> dt.descr
[('name', '<U16'), ('grades', '<f8', (2,))]
fields

Dictionary of named fields defined for this data type, or None.

The dictionary is indexed by keys that are the names of the fields. Each entry in the dictionary is a tuple fully describing the field:

(dtype, offset[, title])

Offset is limited to C int, which is signed and usually 32 bits. If present, the optional title can be any object (if it is a string or unicode then it will also be a key in the fields dictionary, otherwise it’s meta-data). Notice also that the first two elements of the tuple can be passed directly as arguments to the ndarray.getfield and ndarray.setfield methods.

See also

ndarray.getfield, ndarray.setfield

Examples

>>> import numpy as np
>>> dt = np.dtype([('name', np.str_, 16), ('grades', np.float64, (2,))])
>>> print(dt.fields)
{'grades': (dtype(('float64',(2,))), 16), 'name': (dtype('|S16'), 0)}
flags

Bit-flags describing how this data type is to be interpreted.

Bit-masks are in numpy._core.multiarray as the constants ITEM_HASOBJECT, LIST_PICKLE, ITEM_IS_POINTER, NEEDS_INIT, NEEDS_PYAPI, USE_GETITEM, USE_SETITEM. A full explanation of these flags is in C-API documentation; they are largely useful for user-defined data-types.

The following example demonstrates that operations on this particular dtype requires Python C-API.

Examples

>>> import numpy as np
>>> x = np.dtype([('a', np.int32, 8), ('b', np.float64, 6)])
>>> x.flags
16
>>> np._core.multiarray.NEEDS_PYAPI
16
hasobject

Boolean indicating whether this dtype contains any reference-counted objects in any fields or sub-dtypes.

Recall that what is actually in the ndarray memory representing the Python object is the memory address of that object (a pointer). Special handling may be required, and this attribute is useful for distinguishing data types that may contain arbitrary Python objects and data-types that won’t.

isalignedstruct

Boolean indicating whether the dtype is a struct which maintains field alignment. This flag is sticky, so when combining multiple structs together, it is preserved and produces new dtypes which are also aligned.

isbuiltin

Integer indicating how this dtype relates to the built-in dtypes.

Read-only.

0

if this is a structured array type, with fields

1

if this is a dtype compiled into numpy (such as ints, floats etc)

2

if the dtype is for a user-defined numpy type A user-defined type uses the numpy C-API machinery to extend numpy to handle a new array type. See User-defined data-types in the NumPy manual.

Examples

>>> import numpy as np
>>> dt = np.dtype('i2')
>>> dt.isbuiltin
1
>>> dt = np.dtype('f8')
>>> dt.isbuiltin
1
>>> dt = np.dtype([('field1', 'f8')])
>>> dt.isbuiltin
0
isnative

Boolean indicating whether the byte order of this dtype is native to the platform.

itemsize

The element size of this data-type object.

For 18 of the 21 types this number is fixed by the data-type. For the flexible data-types, this number can be anything.

Examples

>>> import numpy as np
>>> arr = np.array([[1, 2], [3, 4]])
>>> arr.dtype
dtype('int64')
>>> arr.itemsize
8

>>> dt = np.dtype([('name', np.str_, 16), ('grades', np.float64, (2,))])
>>> dt.itemsize
80
kind

A character code (one of ‘biufcmMOSUV’) identifying the general kind of data.

b

boolean

i

signed integer

u

unsigned integer

f

floating-point

c

complex floating-point

m

timedelta

M

datetime

O

object

S

(byte-)string

U

Unicode

V

void

Examples

>>> import numpy as np
>>> dt = np.dtype('i4')
>>> dt.kind
'i'
>>> dt = np.dtype('f8')
>>> dt.kind
'f'
>>> dt = np.dtype([('field1', 'f8')])
>>> dt.kind
'V'
metadata

Either None or a readonly dictionary of metadata (mappingproxy).

The metadata field can be set using any dictionary at data-type creation. NumPy currently has no uniform approach to propagating metadata; although some array operations preserve it, there is no guarantee that others will.

Warning

Although used in certain projects, this feature was long undocumented and is not well supported. Some aspects of metadata propagation are expected to change in the future.

Examples

>>> import numpy as np
>>> dt = np.dtype(float, metadata={"key": "value"})
>>> dt.metadata["key"]
'value'
>>> arr = np.array([1, 2, 3], dtype=dt)
>>> arr.dtype.metadata
mappingproxy({'key': 'value'})

Adding arrays with identical datatypes currently preserves the metadata:

>>> (arr + arr).dtype.metadata
mappingproxy({'key': 'value'})

But if the arrays have different dtype metadata, the metadata may be dropped:

>>> dt2 = np.dtype(float, metadata={"key2": "value2"})
>>> arr2 = np.array([3, 2, 1], dtype=dt2)
>>> (arr + arr2).dtype.metadata is None
True  # The metadata field is cleared so None is returned
name

A bit-width name for this data-type.

Un-sized flexible data-type objects do not have this attribute.

Examples

>>> import numpy as np
>>> x = np.dtype(float)
>>> x.name
'float64'
>>> x = np.dtype([('a', np.int32, 8), ('b', np.float64, 6)])
>>> x.name
'void640'
names

Ordered list of field names, or None if there are no fields.

The names are ordered according to increasing byte offset. This can be used, for example, to walk through all of the named fields in offset order.

Examples

>>> dt = np.dtype([('name', np.str_, 16), ('grades', np.float64, (2,))])
>>> dt.names
('name', 'grades')
ndim

Number of dimensions of the sub-array if this data type describes a sub-array, and 0 otherwise.

Examples

>>> import numpy as np
>>> x = np.dtype(float)
>>> x.ndim
0

>>> x = np.dtype((float, 8))
>>> x.ndim
1

>>> x = np.dtype(('i4', (3, 4)))
>>> x.ndim
2
newbyteorder(new_order='S', /)

Return a new dtype with a different byte order.

Changes are also made in all fields and sub-arrays of the data type.

Parameters:

new_order (string, optional) –

Byte order to force; a value from the byte order specifications below. The default value (‘S’) results in swapping the current byte order. new_order codes can be any of:

  • ’S’ - swap dtype from current to opposite endian

  • {‘<’, ‘little’} - little endian

  • {‘>’, ‘big’} - big endian

  • {‘=’, ‘native’} - native order

  • {‘|’, ‘I’} - ignore (no change to byte order)

Returns:

new_dtype – New dtype object with the given change to the byte order.

Return type:

dtype

Notes

Changes are also made in all fields and sub-arrays of the data type.

Examples

>>> import sys
>>> sys_is_le = sys.byteorder == 'little'
>>> native_code = '<' if sys_is_le else '>'
>>> swapped_code = '>' if sys_is_le else '<'
>>> import numpy as np
>>> native_dt = np.dtype(native_code+'i2')
>>> swapped_dt = np.dtype(swapped_code+'i2')
>>> native_dt.newbyteorder('S') == swapped_dt
True
>>> native_dt.newbyteorder() == swapped_dt
True
>>> native_dt == swapped_dt.newbyteorder('S')
True
>>> native_dt == swapped_dt.newbyteorder('=')
True
>>> native_dt == swapped_dt.newbyteorder('N')
True
>>> native_dt == native_dt.newbyteorder('|')
True
>>> np.dtype('<i2') == native_dt.newbyteorder('<')
True
>>> np.dtype('<i2') == native_dt.newbyteorder('L')
True
>>> np.dtype('>i2') == native_dt.newbyteorder('>')
True
>>> np.dtype('>i2') == native_dt.newbyteorder('B')
True
num

A unique number for each of the 21 different built-in types.

These are roughly ordered from least-to-most precision.

Examples

>>> import numpy as np
>>> dt = np.dtype(str)
>>> dt.num
19

>>> dt = np.dtype(float)
>>> dt.num
12
shape

Shape tuple of the sub-array if this data type describes a sub-array, and () otherwise.

Examples

>>> import numpy as np
>>> dt = np.dtype(('i4', 4))
>>> dt.shape
(4,)

>>> dt = np.dtype(('i4', (2, 3)))
>>> dt.shape
(2, 3)
str

The array-protocol typestring of this data-type object.

subdtype

Tuple (item_dtype, shape) if this dtype describes a sub-array, and None otherwise.

The shape is the fixed shape of the sub-array described by this data type, and item_dtype the data type of the array.

If a field whose dtype object has this attribute is retrieved, then the extra dimensions implied by shape are tacked on to the end of the retrieved array.

See also

dtype.base

Examples

>>> import numpy as np
>>> x = numpy.dtype('8f')
>>> x.subdtype
(dtype('float32'), (8,))

>>> x =  numpy.dtype('i2')
>>> x.subdtype
>>>
type = None
quchip.declarative.qnp.ediff1d(ary, to_end=None, to_begin=None)

Compute the differences of the elements of the flattened array.

JAX implementation of numpy.ediff1d().

Args:

ary: input array or scalar. to_end: scalar or array, optional, default=None. Specifies the numbers to

append to the resulting array.

to_begin: scalar or array, optional, default=None. Specifies the numbers to

prepend to the resulting array.

Returns:

An array containing the differences between the elements of the input array.

Note:

Unlike NumPy’s implementation of ediff1d, jax.numpy.ediff1d() will not issue an error if casting to_end or to_begin to the type of ary loses precision.

See also:

  • jax.numpy.diff(): Computes the n-th order difference between elements of the array along a given axis.

  • jax.numpy.cumsum(): Computes the cumulative sum of the elements of the array along a given axis.

  • jax.numpy.gradient(): Computes the gradient of an N-dimensional array.

Examples:
>>> a = jnp.array([2, 3, 5, 9, 1, 4])
>>> jnp.ediff1d(a)
Array([ 1,  2,  4, -8,  3], dtype=int32)
>>> jnp.ediff1d(a, to_begin=-10)
Array([-10,   1,   2,   4,  -8,   3], dtype=int32)
>>> jnp.ediff1d(a, to_end=jnp.array([20, 30]))
Array([ 1,  2,  4, -8,  3, 20, 30], dtype=int32)
>>> jnp.ediff1d(a, to_begin=-10, to_end=jnp.array([20, 30]))
Array([-10,   1,   2,   4,  -8,   3,  20,  30], dtype=int32)

For array with ndim > 1, the differences are computed after flattening the input array.

>>> a1 = jnp.array([[2, -1, 4, 7],
...                 [3, 5, -6, 9]])
>>> jnp.ediff1d(a1)
Array([ -3,   5,   3,  -4,   2, -11,  15], dtype=int32)
>>> a2 = jnp.array([2, -1, 4, 7, 3, 5, -6, 9])
>>> jnp.ediff1d(a2)
Array([ -3,   5,   3,  -4,   2, -11,  15], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.einsum(subscripts, /, *operands, out=None, optimize='auto', precision=None, preferred_element_type=None, _dot_general=<function dot_general>, out_sharding=None)

Einstein summation

JAX implementation of numpy.einsum().

einsum is a powerful and generic API for computing various reductions, inner products, outer products, axis reorderings, and combinations thereof across one or more input arrays. It has a somewhat complicated overloaded API; the arguments below reflect the most common calling convention. The Examples section below demonstrates some of the alternative calling conventions.

Args:

subscripts: string containing axes names separated by commas. *operands: sequence of one or more arrays corresponding to the subscripts. optimize: specify how to optimize the order of computation. In JAX this defaults

to "auto" which produces optimized expressions via the opt_einsum package. Other options are True (same as "optimal"), False (unoptimized), or any string supported by opt_einsum, which includes "optimal", "greedy", "eager", and others. It may also be a pre-computed path (see einsum_path()).

precision: either None (default), which means the default precision for

the backend, a Precision enum value (Precision.DEFAULT, Precision.HIGH or Precision.HIGHEST).

preferred_element_type: either None (default), which means the default

accumulation type for the input types, or a datatype, indicating to accumulate results to and return a result with that datatype.

out: unsupported by JAX _dot_general: optionally override the dot_general callable used by einsum.

This parameter is experimental, and may be removed without warning at any time.

Returns:

array containing the result of the einstein summation.

See also:

jax.numpy.einsum_path()

Examples:

The mechanics of einsum are perhaps best demonstrated by example. Here we show how to use einsum to compute a number of quantities from one or more arrays. For more discussion and examples of einsum, see the documentation of numpy.einsum().

>>> M = jnp.arange(16).reshape(4, 4)
>>> x = jnp.arange(4)
>>> y = jnp.array([5, 4, 3, 2])

Vector product

>>> jnp.einsum('i,i', x, y)
Array(16, dtype=int32)
>>> jnp.vecdot(x, y)
Array(16, dtype=int32)

Here are some alternative einsum calling conventions to compute the same result:

>>> jnp.einsum('i,i->', x, y)  # explicit form
Array(16, dtype=int32)
>>> jnp.einsum(x, (0,), y, (0,))  # implicit form via indices
Array(16, dtype=int32)
>>> jnp.einsum(x, (0,), y, (0,), ())  # explicit form via indices
Array(16, dtype=int32)

Matrix product

>>> jnp.einsum('ij,j->i', M, x)  # explicit form
Array([14, 38, 62, 86], dtype=int32)
>>> jnp.matmul(M, x)
Array([14, 38, 62, 86], dtype=int32)

Here are some alternative einsum calling conventions to compute the same result:

>>> jnp.einsum('ij,j', M, x) # implicit form
Array([14, 38, 62, 86], dtype=int32)
>>> jnp.einsum(M, (0, 1), x, (1,), (0,)) # explicit form via indices
Array([14, 38, 62, 86], dtype=int32)
>>> jnp.einsum(M, (0, 1), x, (1,))  # implicit form via indices
Array([14, 38, 62, 86], dtype=int32)

Outer product

>>> jnp.einsum("i,j->ij", x, y)
Array([[ 0,  0,  0,  0],
       [ 5,  4,  3,  2],
       [10,  8,  6,  4],
       [15, 12,  9,  6]], dtype=int32)
>>> jnp.outer(x, y)
Array([[ 0,  0,  0,  0],
       [ 5,  4,  3,  2],
       [10,  8,  6,  4],
       [15, 12,  9,  6]], dtype=int32)

Some other ways of computing outer products:

>>> jnp.einsum("i,j", x, y)  # implicit form
Array([[ 0,  0,  0,  0],
       [ 5,  4,  3,  2],
       [10,  8,  6,  4],
       [15, 12,  9,  6]], dtype=int32)
>>> jnp.einsum(x, (0,), y, (1,), (0, 1))  # explicit form via indices
Array([[ 0,  0,  0,  0],
       [ 5,  4,  3,  2],
       [10,  8,  6,  4],
       [15, 12,  9,  6]], dtype=int32)
>>> jnp.einsum(x, (0,), y, (1,))  # implicit form via indices
Array([[ 0,  0,  0,  0],
       [ 5,  4,  3,  2],
       [10,  8,  6,  4],
       [15, 12,  9,  6]], dtype=int32)

1D array sum

>>> jnp.einsum("i->", x)  # requires explicit form
Array(6, dtype=int32)
>>> jnp.einsum(x, (0,), ())  # explicit form via indices
Array(6, dtype=int32)
>>> jnp.sum(x)
Array(6, dtype=int32)

Sum along an axis

>>> jnp.einsum("...j->...", M)  # requires explicit form
Array([ 6, 22, 38, 54], dtype=int32)
>>> jnp.einsum(M, (..., 0), (...,))  # explicit form via indices
Array([ 6, 22, 38, 54], dtype=int32)
>>> M.sum(-1)
Array([ 6, 22, 38, 54], dtype=int32)

Matrix transpose

>>> y = jnp.array([[1, 2, 3],
...                [4, 5, 6]])
>>> jnp.einsum("ij->ji", y)  # explicit form
Array([[1, 4],
       [2, 5],
       [3, 6]], dtype=int32)
>>> jnp.einsum("ji", y)  # implicit form
Array([[1, 4],
       [2, 5],
       [3, 6]], dtype=int32)
>>> jnp.einsum(y, (1, 0))  # implicit form via indices
Array([[1, 4],
       [2, 5],
       [3, 6]], dtype=int32)
>>> jnp.einsum(y, (0, 1), (1, 0))  # explicit form via indices
Array([[1, 4],
       [2, 5],
       [3, 6]], dtype=int32)
>>> jnp.transpose(y)
Array([[1, 4],
       [2, 5],
       [3, 6]], dtype=int32)

Matrix diagonal

>>> jnp.einsum("ii->i", M)
Array([ 0,  5, 10, 15], dtype=int32)
>>> jnp.diagonal(M)
Array([ 0,  5, 10, 15], dtype=int32)

Matrix trace

>>> jnp.einsum("ii", M)
Array(30, dtype=int32)
>>> jnp.trace(M)
Array(30, dtype=int32)

Tensor products

>>> x = jnp.arange(30).reshape(2, 3, 5)
>>> y = jnp.arange(60).reshape(3, 4, 5)
>>> jnp.einsum('ijk,jlk->il', x, y)  # explicit form
Array([[ 3340,  3865,  4390,  4915],
       [ 8290,  9940, 11590, 13240]], dtype=int32)
>>> jnp.tensordot(x, y, axes=[(1, 2), (0, 2)])
Array([[ 3340,  3865,  4390,  4915],
       [ 8290,  9940, 11590, 13240]], dtype=int32)
>>> jnp.einsum('ijk,jlk', x, y)  # implicit form
Array([[ 3340,  3865,  4390,  4915],
       [ 8290,  9940, 11590, 13240]], dtype=int32)
>>> jnp.einsum(x, (0, 1, 2), y, (1, 3, 2), (0, 3))  # explicit form via indices
Array([[ 3340,  3865,  4390,  4915],
       [ 8290,  9940, 11590, 13240]], dtype=int32)
>>> jnp.einsum(x, (0, 1, 2), y, (1, 3, 2))  # implicit form via indices
Array([[ 3340,  3865,  4390,  4915],
       [ 8290,  9940, 11590, 13240]], dtype=int32)

Chained dot products

>>> w = jnp.arange(5, 9).reshape(2, 2)
>>> x = jnp.arange(6).reshape(2, 3)
>>> y = jnp.arange(-2, 4).reshape(3, 2)
>>> z = jnp.array([[2, 4, 6], [3, 5, 7]])
>>> jnp.einsum('ij,jk,kl,lm->im', w, x, y, z)
Array([[ 481,  831, 1181],
       [ 651, 1125, 1599]], dtype=int32)
>>> jnp.einsum(w, (0, 1), x, (1, 2), y, (2, 3), z, (3, 4))  # implicit, via indices
Array([[ 481,  831, 1181],
       [ 651, 1125, 1599]], dtype=int32)
>>> w @ x @ y @ z  # direct chain of matmuls
Array([[ 481,  831, 1181],
       [ 651, 1125, 1599]], dtype=int32)
>>> jnp.linalg.multi_dot([w, x, y, z])
Array([[ 481,  831, 1181],
       [ 651, 1125, 1599]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.einsum_path(subscripts, /, *operands, optimize='auto')

Evaluates the optimal contraction path without evaluating the einsum.

JAX implementation of numpy.einsum_path(). This function calls into the opt_einsum package, and makes use of its optimization routines.

Args:

subscripts: string containing axes names separated by commas. *operands: sequence of one or more arrays corresponding to the subscripts. optimize: specify how to optimize the order of computation. In JAX this defaults

to "auto". Other options are True (same as "optimize"), False (unoptimized), or any string supported by opt_einsum, which includes "optimize",, "greedy", "eager", and others.

Returns:

A tuple containing the path that may be passed to einsum(), and a printable object representing this optimal path.

Examples:
>>> key1, key2, key3 = jax.random.split(jax.random.key(0), 3)
>>> x = jax.random.randint(key1, minval=-5, maxval=5, shape=(2, 3))
>>> y = jax.random.randint(key2, minval=-5, maxval=5, shape=(3, 100))
>>> z = jax.random.randint(key3, minval=-5, maxval=5, shape=(100, 5))
>>> path, path_info = jnp.einsum_path("ij,jk,kl", x, y, z, optimize="optimal")
>>> print(path)
[(1, 2), (0, 1)]
>>> print(path_info)
      Complete contraction:  ij,jk,kl->il
            Naive scaling:  4
        Optimized scaling:  3
          Naive FLOP count:  9.000e+3
      Optimized FLOP count:  3.060e+3
      Theoretical speedup:  2.941e+0
      Largest intermediate:  1.500e+1 elements
    --------------------------------------------------------------------------------
    scaling        BLAS                current                             remaining
    --------------------------------------------------------------------------------
      3           GEMM              kl,jk->lj                             ij,lj->il
      3           GEMM              lj,ij->il                                il->il

Use the computed path in einsum():

>>> jnp.einsum("ij,jk,kl", x, y, z, optimize=path)
Array([[-754,  324, -142,   82,   50],
       [ 408,  -50,   87,  -29,    7]], dtype=int32)
Parameters:

optimize (bool | str | list[tuple[int, ...]])

Return type:

tuple[list[tuple[int, …]], Any]

quchip.declarative.qnp.empty(shape, dtype=None, *, device=None)

Create an empty array.

JAX implementation of numpy.empty(). Because XLA cannot create an un-initialized array, jax.numpy.empty() will always return an array full of zeros.

Args:

shape: int or sequence of ints specifying the shape of the created array. dtype: optional dtype for the created array; defaults to float32 or float64

depending on the X64 configuration (see default-dtypes).

device: (optional) Device or Sharding

to which the created array will be committed.

Returns:

Array of the specified shape and dtype, on the specified device if specified.

See also:

  • jax.numpy.empty_like()

  • jax.numpy.zeros()

  • jax.numpy.ones()

  • jax.numpy.full()

Examples:
>>> jnp.empty(4)
Array([0., 0., 0., 0.], dtype=float32)
>>> jnp.empty((2, 3), dtype=bool)
Array([[False, False, False],
       [False, False, False]], dtype=bool)
Parameters:

  • shape (Any)

  • dtype (str | type[Any] | dtype | SupportsDType | None)

  • device (Device | Sharding | None)

Return type:

Array

quchip.declarative.qnp.empty_like(prototype, dtype=None, shape=None, *, device=None)

Create an empty array with the same shape and dtype as an array.

JAX implementation of numpy.empty_like(). Because XLA cannot create an un-initialized array, jax.numpy.empty() will always return an array full of zeros.

Args:

a: Array-like object with shape and dtype attributes. shape: optionally override the shape of the created array. dtype: optionally override the dtype of the created array. device: (optional) Device or Sharding

to which the created array will be committed.

Returns:

Array of the specified shape and dtype, on the specified device if specified.

See also:

  • jax.numpy.empty()

  • jax.numpy.zeros_like()

  • jax.numpy.ones_like()

  • jax.numpy.full_like()

Examples:
>>> x = jnp.arange(4)
>>> jnp.empty_like(x)
Array([0, 0, 0, 0], dtype=int32)
>>> jnp.empty_like(x, dtype=bool)
Array([False, False, False, False], dtype=bool)
>>> jnp.empty_like(x, shape=(2, 3))
Array([[0, 0, 0],
       [0, 0, 0]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.equal(x, y, /)

Returns element-wise truth value of x == y.

JAX implementation of numpy.equal. This function provides the implementation of the == operator for JAX arrays.

Args:

x: input array or scalar. y: input array or scalar. x and y should either have same shape or be

broadcast compatible.

Returns:

A boolean array containing True where the elements of x == y and False otherwise.

See also:

  • jax.numpy.not_equal(): Returns element-wise truth value of x != y.

  • jax.numpy.greater_equal(): Returns element-wise truth value of x >= y.

  • jax.numpy.less_equal(): Returns element-wise truth value of x <= y.

  • jax.numpy.greater(): Returns element-wise truth value of x > y.

  • jax.numpy.less(): Returns element-wise truth value of x < y.

Examples:
>>> jnp.equal(0., -0.)
Array(True, dtype=bool, weak_type=True)
>>> jnp.equal(1, 1.)
Array(True, dtype=bool, weak_type=True)
>>> jnp.equal(5, jnp.array(5))
Array(True, dtype=bool, weak_type=True)
>>> jnp.equal(2, -2)
Array(False, dtype=bool, weak_type=True)
>>> x = jnp.array([[1, 2, 3],
...                [4, 5, 6],
...                [7, 8, 9]])
>>> y = jnp.array([1, 5, 9])
>>> jnp.equal(x, y)
Array([[ True, False, False],
       [False,  True, False],
       [False, False,  True]], dtype=bool)
>>> x == y
Array([[ True, False, False],
       [False,  True, False],
       [False, False,  True]], dtype=bool)
Parameters:
Return type:

Array

quchip.declarative.qnp.exp(x, /)

Calculate element-wise exponential of the input.

JAX implementation of numpy.exp.

Args:

x: input array or scalar

Returns:

An array containing the exponential of each element in x, promotes to inexact dtype.

See also:

  • jax.numpy.log(): Calculates element-wise logarithm of the input.

  • jax.numpy.expm1(): Calculates \(e^x-1\) of each element of the input.

  • jax.numpy.exp2(): Calculates base-2 exponential of each element of the input.

Examples:

jnp.exp follows the properties of exponential such as \(e^{(a+b)} = e^a * e^b\).

>>> x1 = jnp.array([2, 4, 3, 1])
>>> x2 = jnp.array([1, 3, 2, 3])
>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jnp.exp(x1+x2))
[  20.09 1096.63  148.41   54.6 ]
>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jnp.exp(x1)*jnp.exp(x2))
[  20.09 1096.63  148.41   54.6 ]

This property holds for complex input also:

>>> jnp.allclose(jnp.exp(3-4j), jnp.exp(3)*jnp.exp(-4j))
Array(True, dtype=bool)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.exp2(x, /)

Calculate element-wise base-2 exponential of input.

JAX implementation of numpy.exp2.

Args:

x: input array or scalar

Returns:

An array containing the base-2 exponential of each element in x, promotes to inexact dtype.

See also:

  • jax.numpy.log2(): Calculates base-2 logarithm of each element of input.

  • jax.numpy.exp(): Calculates exponential of each element of the input.

  • jax.numpy.expm1(): Calculates \(e^x-1\) of each element of the input.

Examples:

jnp.exp2 follows the properties of the exponential such as \(2^{a+b} = 2^a * 2^b\).

>>> x1 = jnp.array([2, -4, 3, -1])
>>> x2 = jnp.array([-1, 3, -2, 3])
>>> jnp.exp2(x1+x2)
Array([2. , 0.5, 2. , 4. ], dtype=float32)
>>> jnp.exp2(x1)*jnp.exp2(x2)
Array([2. , 0.5, 2. , 4. ], dtype=float32)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.expand_dims(a, axis)

Insert dimensions of length 1 into array

JAX implementation of numpy.expand_dims(), implemented via jax.lax.expand_dims().

Args:

a: input array axis: integer or sequence of integers specifying positions of axes to add.

Returns:

Copy of a with added dimensions.

Notes:

Unlike numpy.expand_dims(), jax.numpy.expand_dims() will return a copy rather than a view of the input array. However, under JIT, the compiler will optimize away such copies when possible, so this doesn’t have performance impacts in practice.

See Also:

  • jax.numpy.squeeze(): inverse of this operation, i.e. remove length-1 dimensions.

  • jax.lax.expand_dims(): XLA version of this functionality.

Examples:
>>> x = jnp.array([1, 2, 3])
>>> x.shape
(3,)

Expand the leading dimension:

>>> jnp.expand_dims(x, 0)
Array([[1, 2, 3]], dtype=int32)
>>> _.shape
(1, 3)

Expand the trailing dimension:

>>> jnp.expand_dims(x, 1)
Array([[1],
       [2],
       [3]], dtype=int32)
>>> _.shape
(3, 1)

Expand multiple dimensions:

>>> jnp.expand_dims(x, (0, 1, 3))
Array([[[[1],
         [2],
         [3]]]], dtype=int32)
>>> _.shape
(1, 1, 3, 1)

Dimensions can also be expanded more succinctly by indexing with None:

>>> x[None]  # equivalent to jnp.expand_dims(x, 0)
Array([[1, 2, 3]], dtype=int32)
>>> x[:, None]  # equivalent to jnp.expand_dims(x, 1)
Array([[1],
       [2],
       [3]], dtype=int32)
>>> x[None, None, :, None]  # equivalent to jnp.expand_dims(x, (0, 1, 3))
Array([[[[1],
         [2],
         [3]]]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.expm1(x, /)

Calculate exp(x)-1 of each element of the input.

JAX implementation of numpy.expm1.

Args:

x: input array or scalar.

Returns:

An array containing exp(x)-1 of each element in x, promotes to inexact dtype.

Note:

jnp.expm1 has much higher precision than the naive computation of exp(x)-1 for small values of x.

See also:

  • jax.numpy.log1p(): Calculates element-wise logarithm of one plus input.

  • jax.numpy.exp(): Calculates element-wise exponential of the input.

  • jax.numpy.exp2(): Calculates base-2 exponential of each element of the input.

Examples:
>>> x = jnp.array([2, -4, 3, -1])
>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jnp.expm1(x))
[ 6.39 -0.98 19.09 -0.63]
>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jnp.exp(x)-1)
[ 6.39 -0.98 19.09 -0.63]

For values very close to 0, jnp.expm1(x) is much more accurate than jnp.exp(x)-1:

>>> x1 = jnp.array([1e-4, 1e-6, 2e-10])
>>> jnp.expm1(x1)
Array([1.0000500e-04, 1.0000005e-06, 2.0000000e-10], dtype=float32)
>>> jnp.exp(x1)-1
Array([1.00016594e-04, 9.53674316e-07, 0.00000000e+00], dtype=float32)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.extract(condition, arr, *, size=None, fill_value=0)

Return the elements of an array that satisfy a condition.

JAX implementation of numpy.extract().

Args:

condition: array of conditions. Will be converted to boolean and flattened to 1D. arr: array of values to extract. Will be flattened to 1D. size: optional static size for output. Must be specified in order for extract

to be compatible with JAX transformations like jit() or vmap().

fill_value: if size is specified, fill padded entries with this value (default: 0).

Returns:

1D array of extracted entries . If size is specified, the result will have shape (size,) and be right-padded with fill_value. If size is not specified, the output shape will depend on the number of True entries in condition.

Notes:

This function does not require strict shape agreement between condition and arr. If condition.size > arr.size, then condition will be truncated, and if arr.size > condition.size, then arr will be truncated.

See also:

jax.numpy.compress(): multi-dimensional version of extract.

Examples:

Extract values from a 1D array:

>>> x = jnp.array([1, 2, 3, 4, 5, 6])
>>> mask = (x % 2 == 0)
>>> jnp.extract(mask, x)
Array([2, 4, 6], dtype=int32)

In the simplest case, this is equivalent to boolean indexing:

>>> x[mask]
Array([2, 4, 6], dtype=int32)

For use with JAX transformations, you can pass the size argument to specify a static shape for the output, along with an optional fill_value that defaults to zero:

>>> jnp.extract(mask, x, size=len(x), fill_value=0)
Array([2, 4, 6, 0, 0, 0], dtype=int32)

Notice that unlike with boolean indexing, extract does not require strict agreement between the sizes of the array and condition, and will effectively truncate both to the minimum size:

>>> short_mask = jnp.array([False, True])
>>> jnp.extract(short_mask, x)
Array([2], dtype=int32)
>>> long_mask = jnp.array([True, False, True, False, False, False, False, False])
>>> jnp.extract(long_mask, x)
Array([1, 3], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.eye(N, M=None, k=0, dtype=None, *, device=None)

Create a square or rectangular identity matrix

JAX implementation of numpy.eye().

Args:

N: integer specifying the first dimension of the array. M: optional integer specifying the second dimension of the array;

defaults to the same value as N.

k: optional integer specifying the offset of the diagonal. Use positive

values for upper diagonals, and negative values for lower diagonals. Default is zero.

dtype: optional dtype; defaults to floating point. device: optional Device or Sharding

to which the created array will be committed.

Returns:

Identity array of shape (N, M), or (N, N) if M is not specified.

See also:

jax.numpy.identity(): Simpler API for generating square identity matrices.

Examples:

A simple 3x3 identity matrix:

>>> jnp.eye(3)
Array([[1., 0., 0.],
       [0., 1., 0.],
       [0., 0., 1.]], dtype=float32)

Integer identity matrices with offset diagonals:

>>> jnp.eye(3, k=1, dtype=int)
Array([[0, 1, 0],
       [0, 0, 1],
       [0, 0, 0]], dtype=int32)
>>> jnp.eye(3, k=-1, dtype=int)
Array([[0, 0, 0],
       [1, 0, 0],
       [0, 1, 0]], dtype=int32)

Non-square identity matrix:

>>> jnp.eye(3, 5, k=1)
Array([[0., 1., 0., 0., 0.],
       [0., 0., 1., 0., 0.],
       [0., 0., 0., 1., 0.]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.fabs(x, /)

Compute the element-wise absolute values of the real-valued input.

JAX implementation of numpy.fabs.

Args:

x: input array or scalar. Must not have a complex dtype.

Returns:

An array with same shape as x and dtype float, containing the element-wise absolute values.

See also:

  • jax.numpy.absolute(): Computes the absolute values of the input including complex dtypes.

  • jax.numpy.abs(): Computes the absolute values of the input including complex dtypes.

Examples:

For integer inputs:

>>> x = jnp.array([-5, -9, 1, 10, 15])
>>> jnp.fabs(x)
Array([ 5.,  9.,  1., 10., 15.], dtype=float32)

For float type inputs:

>>> x1 = jnp.array([-1.342, 5.649, 3.927])
>>> jnp.fabs(x1)
Array([1.342, 5.649, 3.927], dtype=float32)

For boolean inputs:

>>> x2 = jnp.array([True, False])
>>> jnp.fabs(x2)
Array([1., 0.], dtype=float32)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.fill_diagonal(a, val, wrap=False, *, inplace=True)

Return a copy of the array with the diagonal overwritten.

JAX implementation of numpy.fill_diagonal().

The semantics of numpy.fill_diagonal() are to modify arrays in-place, which is not possible for JAX’s immutable arrays. The JAX version returns a modified copy of the input, and adds the inplace parameter which must be set to False` by the user as a reminder of this API difference.

Args:
a: input array. Must have a.ndim >= 2. If a.ndim >= 3, then all

dimensions must be the same size.

val: scalar or array with which to fill the diagonal. If an array, it will

be flattened and repeated to fill the diagonal entries.

inplace: must be set to False to indicate that the input is not modified

in-place, but rather a modified copy is returned.

Returns:

A copy of a with the diagonal set to val.

Examples:
>>> x = jnp.zeros((3, 3), dtype=int)
>>> jnp.fill_diagonal(x, jnp.array([1, 2, 3]), inplace=False)
Array([[1, 0, 0],
       [0, 2, 0],
       [0, 0, 3]], dtype=int32)

Unlike numpy.fill_diagonal(), the input x is not modified.

If the diagonal value has too many entries, it will be truncated

>>> jnp.fill_diagonal(x, jnp.arange(100, 200), inplace=False)
Array([[100,   0,   0],
       [  0, 101,   0],
       [  0,   0, 102]], dtype=int32)

If the diagonal has too few entries, it will be repeated:

>>> x = jnp.zeros((4, 4), dtype=int)
>>> jnp.fill_diagonal(x, jnp.array([3, 4]), inplace=False)
Array([[3, 0, 0, 0],
       [0, 4, 0, 0],
       [0, 0, 3, 0],
       [0, 0, 0, 4]], dtype=int32)

For non-square arrays, the diagonal of the leading square slice is filled:

>>> x = jnp.zeros((3, 5), dtype=int)
>>> jnp.fill_diagonal(x, 1, inplace=False)
Array([[1, 0, 0, 0, 0],
       [0, 1, 0, 0, 0],
       [0, 0, 1, 0, 0]], dtype=int32)

And for square N-dimensional arrays, the N-dimensional diagonal is filled:

>>> y = jnp.zeros((2, 2, 2))
>>> jnp.fill_diagonal(y, 1, inplace=False)
Array([[[1., 0.],
        [0., 0.]],

       [[0., 0.],
        [0., 1.]]], dtype=float32)
Parameters:
Return type:

Array

class quchip.declarative.qnp.finfo(dtype)[source]

Bases: finfo

Machine limits for floating point types.

bits

The number of bits occupied by the type.

Type:

int

dtype

Returns the dtype for which finfo returns information. For complex input, the returned dtype is the associated float* dtype for its real and complex components.

Type:

dtype

eps

The difference between 1.0 and the next smallest representable float larger than 1.0. For example, for 64-bit binary floats in the IEEE-754 standard, eps = 2**-52, approximately 2.22e-16.

Type:

float

epsneg

The difference between 1.0 and the next smallest representable float less than 1.0. For example, for 64-bit binary floats in the IEEE-754 standard, epsneg = 2**-53, approximately 1.11e-16.

Type:

float

iexp

The number of bits in the exponent portion of the floating point representation.

Type:

int

machep

The exponent that yields eps.

Type:

int

max

The largest representable number.

Type:

floating point number of the appropriate type

maxexp

The smallest positive power of the base (2) that causes overflow.

Type:

int

min

The smallest representable number, typically -max.

Type:

floating point number of the appropriate type

minexp

The most negative power of the base (2) consistent with there being no leading 0’s in the mantissa.

Type:

int

negep

The exponent that yields epsneg.

Type:

int

nexp

The number of bits in the exponent including its sign and bias.

Type:

int

nmant

The number of bits in the mantissa.

Type:

int

precision

The approximate number of decimal digits to which this kind of float is precise.

Type:

int

resolution

The approximate decimal resolution of this type, i.e., 10**-precision.

Type:

floating point number of the appropriate type

tiny

An alias for smallest_normal, kept for backwards compatibility.

Type:

float

smallest_normal

The smallest positive floating point number with 1 as leading bit in the mantissa following IEEE-754 (see Notes).

Type:

float

smallest_subnormal

The smallest positive floating point number with 0 as leading bit in the mantissa following IEEE-754.

Type:

float

Parameters:

dtype (float, dtype, or instance) – Kind of floating point or complex floating point data-type about which to get information.

See also

iinfo

The equivalent for integer data types.

spacing

The distance between a value and the nearest adjacent number

nextafter

The next floating point value after x1 towards x2

Notes

For developers of NumPy: do not instantiate this at the module level. The initial calculation of these parameters is expensive and negatively impacts import times. These objects are cached, so calling finfo() repeatedly inside your functions is not a problem.

Note that smallest_normal is not actually the smallest positive representable value in a NumPy floating point type. As in the IEEE-754 standard [1], NumPy floating point types make use of subnormal numbers to fill the gap between 0 and smallest_normal. However, subnormal numbers may have significantly reduced precision [2].

This function can also be used for complex data types as well. If used, the output will be the same as the corresponding real float type (e.g. numpy.finfo(numpy.csingle) is the same as numpy.finfo(numpy.single)). However, the output is true for the real and imaginary components.

References

Examples

>>> import numpy as np
>>> np.finfo(np.float64).dtype
dtype('float64')
>>> np.finfo(np.complex64).dtype
dtype('float32')
quchip.declarative.qnp.fix(x, out=None)

Round input to the nearest integer towards zero.

JAX implementation of numpy.fix().

Args:

x: input array. out: unused by JAX.

Returns:

An array with same shape and dtype as x containing the rounded values.

See also:

  • jax.numpy.trunc(): Rounds the input to nearest integer towards zero.

  • jax.numpy.ceil(): Rounds the input up to the nearest integer.

  • jax.numpy.floor(): Rounds the input down to the nearest integer.

Examples:
>>> key = jax.random.key(0)
>>> x = jax.random.uniform(key, (3, 3), minval=-5, maxval=5)
>>> with jnp.printoptions(precision=2, suppress=True):
...     print(x)
[[ 4.48  4.79 -1.68]
 [-0.31  0.7  -3.34]
 [-1.9   1.89  2.47]]
>>> jnp.fix(x)
Array([[ 4.,  4., -1.],
       [-0.,  0., -3.],
       [-1.,  1.,  2.]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.flatnonzero(a, *, size=None, fill_value=None)

Return indices of nonzero elements in a flattened array

JAX implementation of numpy.flatnonzero().

jnp.flatnonzero(x) is equivalent to nonzero(ravel(a))[0]. For a full discussion of the parameters to this function, refer to jax.numpy.nonzero().

Args:

a: N-dimensional array. size: optional static integer specifying the number of nonzero entries to

return. See jax.numpy.nonzero() for more discussion of this parameter.

fill_value: optional padding value when size is specified. Defaults to 0.

See jax.numpy.nonzero() for more discussion of this parameter.

Returns:

Array containing the indices of each nonzero value in the flattened array.

See Also:

  • jax.numpy.nonzero()

  • jax.numpy.where()

Examples:
>>> x = jnp.array([[0, 5, 0],
...                [6, 0, 8]])
>>> jnp.flatnonzero(x)
Array([1, 3, 5], dtype=int32)

This is equivalent to calling nonzero() on the flattened array, and extracting the first entry in the resulting tuple:

>>> jnp.nonzero(x.ravel())[0]
Array([1, 3, 5], dtype=int32)

The returned indices can be used to extract nonzero entries from the flattened array:

>>> indices = jnp.flatnonzero(x)
>>> x.ravel()[indices]
Array([5, 6, 8], dtype=int32)
Parameters:
Return type:

Array

class quchip.declarative.qnp.flexible

Bases: generic

Abstract base class of all scalar types without predefined length. The actual size of these types depends on the specific numpy.dtype instantiation.

quchip.declarative.qnp.flip(m, axis=None)

Reverse the order of elements of an array along the given axis.

JAX implementation of numpy.flip().

Args:

m: Array. axis: integer or sequence of integers. Specifies along which axis or axes

should the array elements be reversed. Default is None, which flips along all axes.

Returns:

An array with the elements in reverse order along axis.

See Also:

  • jax.numpy.fliplr(): reverse the order along axis 1 (left/right)

  • jax.numpy.flipud(): reverse the order along axis 0 (up/down)

Examples:
>>> x1 = jnp.array([[1, 2],
...                 [3, 4]])
>>> jnp.flip(x1)
Array([[4, 3],
       [2, 1]], dtype=int32)

If axis is specified with an integer, then jax.numpy.flip reverses the array along that particular axis only.

>>> jnp.flip(x1, axis=1)
Array([[2, 1],
       [4, 3]], dtype=int32)

>>> x2 = jnp.arange(1, 9).reshape(2, 2, 2)
>>> x2
Array([[[1, 2],
        [3, 4]],

       [[5, 6],
        [7, 8]]], dtype=int32)
>>> jnp.flip(x2)
Array([[[8, 7],
        [6, 5]],

       [[4, 3],
        [2, 1]]], dtype=int32)

When axis is specified with a sequence of integers, then jax.numpy.flip reverses the array along the specified axes.

>>> jnp.flip(x2, axis=[1, 2])
Array([[[4, 3],
        [2, 1]],

       [[8, 7],
        [6, 5]]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.fliplr(m)

Reverse the order of elements of an array along axis 1.

JAX implementation of numpy.fliplr().

Args:

m: Array with at least two dimensions.

Returns:

An array with the elements in reverse order along axis 1.

See Also:

  • jax.numpy.flip(): reverse the order along the given axis

  • jax.numpy.flipud(): reverse the order along axis 0

Examples:
>>> x = jnp.array([[1, 2],
...                [3, 4]])
>>> jnp.fliplr(x)
Array([[2, 1],
       [4, 3]], dtype=int32)
Parameters:

m (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.flipud(m)

Reverse the order of elements of an array along axis 0.

JAX implementation of numpy.flipud().

Args:

m: Array with at least one dimension.

Returns:

An array with the elements in reverse order along axis 0.

See Also:

  • jax.numpy.flip(): reverse the order along the given axis

  • jax.numpy.fliplr(): reverse the order along axis 1

Examples:
>>> x = jnp.array([[1, 2],
...                [3, 4]])
>>> jnp.flipud(x)
Array([[3, 4],
       [1, 2]], dtype=int32)
Parameters:

m (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

class quchip.declarative.qnp.float16(x)

Bases: object

A JAX scalar constructor of type float16.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype('float16')
class quchip.declarative.qnp.float32(x)

Bases: object

A JAX scalar constructor of type float32.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype('float32')
class quchip.declarative.qnp.float4_e2m1fn(x)

Bases: object

A JAX scalar constructor of type float4_e2m1fn.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype(float4_e2m1fn)
class quchip.declarative.qnp.float64(x)

Bases: object

A JAX scalar constructor of type float64.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype('float64')
class quchip.declarative.qnp.float8_e3m4(x)

Bases: object

A JAX scalar constructor of type float8_e3m4.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype(float8_e3m4)
class quchip.declarative.qnp.float8_e4m3(x)

Bases: object

A JAX scalar constructor of type float8_e4m3.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype(float8_e4m3)
class quchip.declarative.qnp.float8_e4m3b11fnuz(x)

Bases: object

A JAX scalar constructor of type float8_e4m3b11fnuz.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype(float8_e4m3b11fnuz)
class quchip.declarative.qnp.float8_e4m3fn(x)

Bases: object

A JAX scalar constructor of type float8_e4m3fn.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype(float8_e4m3fn)
class quchip.declarative.qnp.float8_e4m3fnuz(x)

Bases: object

A JAX scalar constructor of type float8_e4m3fnuz.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype(float8_e4m3fnuz)
class quchip.declarative.qnp.float8_e5m2(x)

Bases: object

A JAX scalar constructor of type float8_e5m2.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype(float8_e5m2)
class quchip.declarative.qnp.float8_e5m2fnuz(x)

Bases: object

A JAX scalar constructor of type float8_e5m2fnuz.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype(float8_e5m2fnuz)
class quchip.declarative.qnp.float8_e8m0fnu(x)

Bases: object

A JAX scalar constructor of type float8_e8m0fnu.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype(float8_e8m0fnu)
quchip.declarative.qnp.float_

alias of float64

quchip.declarative.qnp.float_power(x, y, /)

Calculate element-wise base x exponential of y.

JAX implementation of numpy.float_power.

Args:

x: scalar or array. Specifies the bases. y: scalar or array. Specifies the exponents. x and y should either

have same shape or be broadcast compatible.

Returns:

An array containing the base x exponentials of y, promoting to the inexact dtype.

See also:

  • jax.numpy.exp(): Calculates element-wise exponential of the input.

  • jax.numpy.exp2(): Calculates base-2 exponential of each element of the input.

Examples:

Inputs with same shape:

>>> x = jnp.array([3, 1, -5])
>>> y = jnp.array([2, 4, -1])
>>> jnp.float_power(x, y)
Array([ 9. ,  1. , -0.2], dtype=float32)

Inputs with broadcast compatibility:

>>> x1 = jnp.array([[2, -4, 1],
...                 [-1, 2, 3]])
>>> y1 = jnp.array([-2, 1, 4])
>>> jnp.float_power(x1, y1)
Array([[ 0.25, -4.  ,  1.  ],
       [ 1.  ,  2.  , 81.  ]], dtype=float32)

jnp.float_power produces nan for negative values raised to a non-integer values.

>>> jnp.float_power(-3, 1.7)
Array(nan, dtype=float32, weak_type=True)
Parameters:
Return type:

Array

class quchip.declarative.qnp.floating

Bases: inexact

Abstract base class of all floating-point scalar types.

quchip.declarative.qnp.floor(x, /)

Round input to the nearest integer downwards.

JAX implementation of numpy.floor.

Args:

x: input array or scalar. Must not have complex dtype.

Returns:

An array with same shape and dtype as x containing the values rounded to the nearest integer that is less than or equal to the value itself.

See also:

  • jax.numpy.fix(): Rounds the input to the nearest integer towards zero.

  • jax.numpy.trunc(): Rounds the input to the nearest integer towards zero.

  • jax.numpy.ceil(): Rounds the input up to the nearest integer.

Examples:
>>> key = jax.random.key(42)
>>> x = jax.random.uniform(key, (3, 3), minval=-5, maxval=5)
>>> with jnp.printoptions(precision=2, suppress=True):
...     print(x)
[[-0.11  1.8   1.16]
 [ 0.61 -0.49  0.86]
 [-4.25  2.75  1.99]]
>>> jnp.floor(x)
Array([[-1.,  1.,  1.],
       [ 0., -1.,  0.],
       [-5.,  2.,  1.]], dtype=float32)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.floor_divide(x1, x2, /)

Calculates the floor division of x1 by x2 element-wise

JAX implementation of numpy.floor_divide.

Args:

x1: Input array, the dividend x2: Input array, the divisor

Returns:

An array-like object containing each of the quotients rounded down to the nearest integer towards negative infinity. This is equivalent to x1 // x2 in Python.

Note:

x1 // x2 is equivalent to jnp.floor_divide(x1, x2) for arrays x1 and x2

See Also:

jax.numpy.divide() and jax.numpy.true_divide() for floating point division.

Examples:
>>> x1 = jnp.array([10, 20, 30])
>>> x2 = jnp.array([3, 4, 7])
>>> jnp.floor_divide(x1, x2)
Array([3, 5, 4], dtype=int32)

>>> x1 = jnp.array([-5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5])
>>> x2 = 3
>>> jnp.floor_divide(x1, x2)
Array([-2, -2, -1, -1, -1,  0,  0,  0,  1,  1,  1], dtype=int32)

>>> x1 = jnp.array([6, 6, 6], dtype=jnp.int32)
>>> x2 = jnp.array([2.0, 2.5, 3.0], dtype=jnp.float32)
>>> jnp.floor_divide(x1, x2)
Array([3., 2., 2.], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.fmax(x1, x2)

Return element-wise maximum of the input arrays.

JAX implementation of numpy.fmax().

Args:

x1: input array or scalar x2: input array or scalar. x1 and x1 must either have same shape or be

broadcast compatible.

Returns:

An array containing the element-wise maximum of x1 and x2.

Note:
For each pair of elements, jnp.fmax returns:

  • the larger of the two if both elements are finite numbers.

  • finite number if one element is nan.

  • nan if both elements are nan.

  • inf if one element is inf and the other is finite or nan.

  • -inf if one element is -inf and the other is nan.

Examples:
>>> jnp.fmax(3, 7)
Array(7, dtype=int32, weak_type=True)
>>> jnp.fmax(5, jnp.array([1, 7, 9, 4]))
Array([5, 7, 9, 5], dtype=int32)

>>> x1 = jnp.array([1, 3, 7, 8])
>>> x2 = jnp.array([-1, 4, 6, 9])
>>> jnp.fmax(x1, x2)
Array([1, 4, 7, 9], dtype=int32)

>>> x3 = jnp.array([[2, 3, 5, 10],
...                 [11, 9, 7, 5]])
>>> jnp.fmax(x1, x3)
Array([[ 2,  3,  7, 10],
       [11,  9,  7,  8]], dtype=int32)

>>> x4 = jnp.array([jnp.inf, 6, -jnp.inf, nan])
>>> x5 = jnp.array([[3, 5, 7, nan],
...                 [nan, 9, nan, -1]])
>>> jnp.fmax(x4, x5)
Array([[ inf,   6.,   7.,  nan],
       [ inf,   9., -inf,  -1.]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.fmin(x1, x2)

Return element-wise minimum of the input arrays.

JAX implementation of numpy.fmin().

Args:

x1: input array or scalar. x2: input array or scalar. x1 and x2 must either have same shape or be

broadcast compatible.

Returns:

An array containing the element-wise minimum of x1 and x2.

Note:
For each pair of elements, jnp.fmin returns:

  • the smaller of the two if both elements are finite numbers.

  • finite number if one element is nan.

  • -inf if one element is -inf and the other is finite or nan.

  • inf if one element is inf and the other is nan.

  • nan if both elements are nan.

Examples:
>>> jnp.fmin(2, 3)
Array(2, dtype=int32, weak_type=True)
>>> jnp.fmin(2, jnp.array([1, 4, 2, -1]))
Array([ 1,  2,  2, -1], dtype=int32)

>>> x1 = jnp.array([1, 3, 2])
>>> x2 = jnp.array([2, 1, 4])
>>> jnp.fmin(x1, x2)
Array([1, 1, 2], dtype=int32)

>>> x3 = jnp.array([1, 5, 3])
>>> x4 = jnp.array([[2, 3, 1],
...                 [5, 6, 7]])
>>> jnp.fmin(x3, x4)
Array([[1, 3, 1],
       [1, 5, 3]], dtype=int32)

>>> nan = jnp.nan
>>> x5 = jnp.array([jnp.inf, 5, nan])
>>> x6 = jnp.array([[2, 3, nan],
...                 [nan, 6, 7]])
>>> jnp.fmin(x5, x6)
Array([[ 2.,  3., nan],
       [inf,  5.,  7.]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.fmod(x1, x2, /)

Calculate element-wise floating-point modulo operation.

JAX implementation of numpy.fmod.

Args:

x1: scalar or array. Specifies the dividend. x2: scalar or array. Specifies the divisor. x1 and x2 should either

have same shape or be broadcast compatible.

Returns:

An array containing the result of the element-wise floating-point modulo operation of x1 and x2 with same sign as the elements of x1.

Note:

The result of jnp.fmod is equivalent to x1 - x2 * jnp.fix(x1 / x2).

See also:

  • jax.numpy.mod() and jax.numpy.remainder(): Returns the element-wise remainder of the division.

  • jax.numpy.divmod(): Calculates the integer quotient and remainder of x1 by x2, element-wise.

Examples:
>>> x1 = jnp.array([[3, -1, 4],
...                 [8, 5, -2]])
>>> x2 = jnp.array([2, 3, -5])
>>> jnp.fmod(x1, x2)
Array([[ 1, -1,  4],
       [ 0,  2, -2]], dtype=int32)
>>> x1 - x2 * jnp.fix(x1 / x2)
Array([[ 1., -1.,  4.],
       [ 0.,  2., -2.]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.frexp(x, /)

Split floating point values into mantissa and twos exponent.

JAX implementation of numpy.frexp().

Args:

x: real-valued array

Returns:

A tuple (mantissa, exponent) where mantissa is a floating point value between -1 and 1, and exponent is an integer such that x == mantissa * 2 ** exponent.

See also:

  • jax.numpy.ldexp(): compute the inverse of frexp.

Examples:

Split values into mantissa and exponent:

>>> x = jnp.array([1., 2., 3., 4., 5.])
>>> m, e = jnp.frexp(x)
>>> m
Array([0.5  , 0.5  , 0.75 , 0.5  , 0.625], dtype=float32)
>>> e
Array([1, 2, 2, 3, 3], dtype=int32)

Reconstruct the original array:

>>> m * 2 ** e
Array([1., 2., 3., 4., 5.], dtype=float32)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

tuple[Array, Array]

quchip.declarative.qnp.from_dlpack(x, /, *, device=None, copy=None)

Construct a JAX array via DLPack.

JAX implementation of numpy.from_dlpack().

Args:
x: An object that implements the DLPack protocol via the __dlpack__

and __dlpack_device__ methods, or a legacy DLPack tensor on either CPU or GPU.

device: An optional Device or Sharding,

representing the single device onto which the returned array should be placed. If given, then the result is committed to the device. If unspecified, the resulting array will be unpacked onto the same device it originated from. Setting device to a device different from the source of external_array will require a copy, meaning copy must be set to either True or None.

copy: An optional boolean, controlling whether or not a copy is performed.

If copy=True then a copy is always performed, even if unpacked onto the same device. If copy=False then the copy is never performed and will raise an error if necessary. When copy=None (default) then a copy may be performed if needed for a device transfer.

Returns:

A JAX array of the input buffer.

Note:

While JAX arrays are always immutable, dlpack buffers cannot be marked as immutable, and it is possible for processes external to JAX to mutate them in-place. If a JAX Array is constructed from a dlpack buffer without copying and the source buffer is later modified in-place, it may lead to undefined behavior when using the associated JAX array.

Examples:

Passing data between NumPy and JAX via DLPack:

>>> import numpy as np
>>> rng = np.random.default_rng(42)
>>> x_numpy = rng.random(4, dtype='float32')
>>> print(x_numpy)
[0.08925092 0.773956   0.6545715  0.43887842]
>>> hasattr(x_numpy, "__dlpack__")  # NumPy supports the DLPack interface
True

>>> import jax.numpy as jnp
>>> x_jax = jnp.from_dlpack(x_numpy)
>>> print(x_jax)
[0.08925092 0.773956   0.6545715  0.43887842]
>>> hasattr(x_jax, "__dlpack__")  # JAX supports the DLPack interface
True

>>> x_numpy_round_trip = np.from_dlpack(x_jax)
>>> print(x_numpy_round_trip)
[0.08925092 0.773956   0.6545715  0.43887842]
Parameters:

  • x (Any)

  • device (Device | Sharding | None)

  • copy (bool | None)

Return type:

Array

quchip.declarative.qnp.frombuffer(buffer, dtype=<class 'float'>, count=-1, offset=0)

Convert a buffer into a 1-D JAX array.

JAX implementation of numpy.frombuffer().

Args:
buffer: an object containing the data. It must be either a bytes object with

a length that is an integer multiple of the dtype element size, or it must be an object exporting the Python buffer interface.

dtype: optional. Desired data type for the array. Default is float64.

This specifies the dtype used to parse the buffer, but note that after parsing, 64-bit values will be cast to 32-bit JAX arrays if the jax_enable_x64 flag is set to False.

count: optional integer specifying the number of items to read from the buffer.

If -1 (default), all items from the buffer are read.

offset: optional integer specifying the number of bytes to skip at the beginning

of the buffer. Default is 0.

Returns:

A 1-D JAX array representing the interpreted data from the buffer.

See also:

  • jax.numpy.fromstring(): convert a string of text into 1-D JAX array.

Examples:

Using a bytes buffer:

>>> buf = b"\x00\x01\x02\x03\x04"
>>> jnp.frombuffer(buf, dtype=jnp.uint8)
Array([0, 1, 2, 3, 4], dtype=uint8)
>>> jnp.frombuffer(buf, dtype=jnp.uint8, offset=1)
Array([1, 2, 3, 4], dtype=uint8)

Constructing a JAX array via the Python buffer interface, using Python’s built-in array module.

>>> from array import array
>>> pybuffer = array('i', [0, 1, 2, 3, 4])
>>> jnp.frombuffer(pybuffer, dtype=jnp.int32)
Array([0, 1, 2, 3, 4], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.fromfile(*args, **kwargs)

Unimplemented JAX wrapper for jnp.fromfile.

This function is left deliberately unimplemented because it may be non-pure and thus unsafe for use with JIT and other JAX transformations. Consider using jnp.asarray(np.fromfile(...)) instead, although care should be taken if np.fromfile is used within jax transformations because of its potential side-effect of consuming the file object; for more information see Common Gotchas: Pure Functions.

quchip.declarative.qnp.fromfunction(function, shape, *, dtype=<class 'float'>, **kwargs)

Create an array from a function applied over indices.

JAX implementation of numpy.fromfunction(). The JAX implementation differs in that it dispatches via jax.vmap(), and so unlike in NumPy the function logically operates on scalar inputs, and need not explicitly handle broadcasted inputs (See Examples below).

Args:

function: a function that takes N dynamic scalars and outputs a scalar. shape: a length-N tuple of integers specifying the output shape. dtype: optionally specify the dtype of the inputs. Defaults to floating-point. kwargs: additional keyword arguments are passed statically to function.

Returns:

An array of shape shape if function returns a scalar, or in general a pytree of arrays with leading dimensions shape, as determined by the output of function.

See also:

  • jax.vmap(): the core transformation that the fromfunction() API is built on.

Examples:

Generate a multiplication table of a given shape:

>>> jnp.fromfunction(jnp.multiply, shape=(3, 6), dtype=int)
Array([[ 0,  0,  0,  0,  0,  0],
       [ 0,  1,  2,  3,  4,  5],
       [ 0,  2,  4,  6,  8, 10]], dtype=int32)

When function returns a non-scalar the output will have leading dimension of shape:

>>> def f(x):
...   return (x + 1) * jnp.arange(3)
>>> jnp.fromfunction(f, shape=(2,))
Array([[0., 1., 2.],
       [0., 2., 4.]], dtype=float32)

function may return multiple results, in which case each is mapped independently:

>>> def f(x, y):
...   return x + y, x * y
>>> x_plus_y, x_times_y = jnp.fromfunction(f, shape=(3, 5))
>>> print(x_plus_y)
[[0. 1. 2. 3. 4.]
 [1. 2. 3. 4. 5.]
 [2. 3. 4. 5. 6.]]
>>> print(x_times_y)
[[0. 0. 0. 0. 0.]
 [0. 1. 2. 3. 4.]
 [0. 2. 4. 6. 8.]]

The JAX implementation differs slightly from NumPy’s implementation. In numpy.fromfunction(), the function is expected to explicitly operate element-wise on the full grid of input values:

>>> def f(x, y):
...   print(f"{x.shape = }\n{y.shape = }")
...   return x + y
...
>>> np.fromfunction(f, (2, 3))
x.shape = (2, 3)
y.shape = (2, 3)
array([[0., 1., 2.],
       [1., 2., 3.]])

In jax.numpy.fromfunction(), the function is vectorized via jax.vmap(), and so is expected to operate on scalar values:

>>> jnp.fromfunction(f, (2, 3))
x.shape = ()
y.shape = ()
Array([[0., 1., 2.],
       [1., 2., 3.]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.fromiter(*args, **kwargs)

Unimplemented JAX wrapper for jnp.fromiter.

This function is left deliberately unimplemented because it may be non-pure and thus unsafe for use with JIT and other JAX transformations. Consider using jnp.asarray(np.fromiter(...)) instead, although care should be taken if np.fromiter is used within jax transformations because of its potential side-effect of consuming the iterable object; for more information see Common Gotchas: Pure Functions.

quchip.declarative.qnp.frompyfunc(func, /, nin, nout, *, identity=None)

Create a JAX ufunc from an arbitrary JAX-compatible scalar function.

Args:

func : a callable that takes nin scalar arguments and returns nout outputs. nin: integer specifying the number of scalar inputs nout: integer specifying the number of scalar outputs identity: (optional) a scalar specifying the identity of the operation, if any.

Returns:

wrapped : jax.numpy.ufunc wrapper of func.

Examples:

Here is an example of creating a ufunc similar to jax.numpy.add:

>>> import operator
>>> add = frompyfunc(operator.add, nin=2, nout=1, identity=0)

Now all the standard jax.numpy.ufunc methods are available:

>>> x = jnp.arange(4)
>>> add(x, 10)
Array([10, 11, 12, 13], dtype=int32)
>>> add.outer(x, x)
Array([[0, 1, 2, 3],
       [1, 2, 3, 4],
       [2, 3, 4, 5],
       [3, 4, 5, 6]], dtype=int32)
>>> add.reduce(x)
Array(6, dtype=int32)
>>> add.accumulate(x)
Array([0, 1, 3, 6], dtype=int32)
>>> add.at(x, 1, 10, inplace=False)
Array([ 0, 11,  2,  3], dtype=int32)
Parameters:
Return type:

ufunc

quchip.declarative.qnp.fromstring(string, dtype=<class 'float'>, count=-1, *, sep)

Convert a string of text into 1-D JAX array.

JAX implementation of numpy.fromstring().

Args:

string: input string containing the data. dtype: optional. Desired data type for the array. Default is float. count: optional integer specifying the number of items to read from the string.

If -1 (default), all items are read.

sep: the string used to separate values in the input string.

Returns:

A 1-D JAX array containing the parsed data from the input string.

See also:

  • jax.numpy.frombuffer(): construct a JAX array from an object that implements the buffer interface.

Examples:
>>> jnp.fromstring("1 2 3", dtype=int, sep=" ")
Array([1, 2, 3], dtype=int32)
>>> jnp.fromstring("0.1, 0.2, 0.3", dtype=float, count=2, sep=",")
Array([0.1, 0.2], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.full(shape, fill_value, dtype=None, *, device=None)

Create an array full of a specified value.

JAX implementation of numpy.full().

Args:

shape: int or sequence of ints specifying the shape of the created array. fill_value: scalar or array with which to fill the created array. dtype: optional dtype for the created array; defaults to the dtype of the

fill value.

device: (optional) Device or Sharding

to which the created array will be committed.

Returns:

Array of the specified shape and dtype, on the specified device if specified.

See also:

  • jax.numpy.full_like()

  • jax.numpy.empty()

  • jax.numpy.zeros()

  • jax.numpy.ones()

Examples:
>>> jnp.full(4, 2, dtype=float)
Array([2., 2., 2., 2.], dtype=float32)
>>> jnp.full((2, 3), 0, dtype=bool)
Array([[False, False, False],
       [False, False, False]], dtype=bool)

fill_value may also be an array that is broadcast to the specified shape:

>>> jnp.full((2, 3), fill_value=jnp.arange(3))
Array([[0, 1, 2],
       [0, 1, 2]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.full_like(a, fill_value, dtype=None, shape=None, *, device=None)

Create an array full of a specified value with the same shape and dtype as an array.

JAX implementation of numpy.full_like().

Args:

a: Array-like object with shape and dtype attributes. fill_value: scalar or array with which to fill the created array. shape: optionally override the shape of the created array. dtype: optionally override the dtype of the created array. device: (optional) Device or Sharding

to which the created array will be committed.

Returns:

Array of the specified shape and dtype, on the specified device if specified.

See also:

  • jax.numpy.full()

  • jax.numpy.empty_like()

  • jax.numpy.zeros_like()

  • jax.numpy.ones_like()

Examples:
>>> x = jnp.arange(4.0)
>>> jnp.full_like(x, 2)
Array([2., 2., 2., 2.], dtype=float32)
>>> jnp.full_like(x, 0, shape=(2, 3))
Array([[0., 0., 0.],
       [0., 0., 0.]], dtype=float32)

fill_value may also be an array that is broadcast to the specified shape:

>>> x = jnp.arange(6).reshape(2, 3)
>>> jnp.full_like(x, fill_value=jnp.array([[1], [2]]))
Array([[1, 1, 1],
       [2, 2, 2]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.gcd(x1, x2)

Compute the greatest common divisor of two arrays.

JAX implementation of numpy.gcd().

Args:

x1: First input array. The elements must have integer dtype. x2: Second input array. The elements must have integer dtype.

Returns:

An array containing the greatest common divisors of the corresponding elements from the absolute values of x1 and x2.

See also:

  • jax.numpy.lcm(): compute the least common multiple of two arrays.

Examples:

Scalar inputs:

>>> jnp.gcd(12, 18)
Array(6, dtype=int32, weak_type=True)

Array inputs:

>>> x1 = jnp.array([12, 18, 24])
>>> x2 = jnp.array([5, 10, 15])
>>> jnp.gcd(x1, x2)
Array([1, 2, 3], dtype=int32)

Broadcasting:

>>> x1 = jnp.array([12])
>>> x2 = jnp.array([6, 9, 12])
>>> jnp.gcd(x1, x2)
Array([ 6,  3, 12], dtype=int32)
Parameters:
Return type:

Array

class quchip.declarative.qnp.generic

Bases: object

Base class for numpy scalar types.

Class from which most (all?) numpy scalar types are derived. For consistency, exposes the same API as ndarray, despite many consequent attributes being either “get-only,” or completely irrelevant. This is the class from which it is strongly suggested users should derive custom scalar types.

T

Scalar attribute identical to the corresponding array attribute.

Please see ndarray.T.

all()

Scalar method identical to the corresponding array attribute.

Please see ndarray.all.

any()

Scalar method identical to the corresponding array attribute.

Please see ndarray.any.

argmax()

Scalar method identical to the corresponding array attribute.

Please see ndarray.argmax.

argmin()

Scalar method identical to the corresponding array attribute.

Please see ndarray.argmin.

argsort()

Scalar method identical to the corresponding array attribute.

Please see ndarray.argsort.

astype()

Scalar method identical to the corresponding array attribute.

Please see ndarray.astype.

base

Scalar attribute identical to the corresponding array attribute.

Please see ndarray.base.

byteswap()

Scalar method identical to the corresponding array attribute.

Please see ndarray.byteswap.

choose()

Scalar method identical to the corresponding array attribute.

Please see ndarray.choose.

clip()

Scalar method identical to the corresponding array attribute.

Please see ndarray.clip.

compress()

Scalar method identical to the corresponding array attribute.

Please see ndarray.compress.

conj()
conjugate()

Scalar method identical to the corresponding array attribute.

Please see ndarray.conjugate.

copy()

Scalar method identical to the corresponding array attribute.

Please see ndarray.copy.

cumprod()

Scalar method identical to the corresponding array attribute.

Please see ndarray.cumprod.

cumsum()

Scalar method identical to the corresponding array attribute.

Please see ndarray.cumsum.

data

Pointer to start of data.

device
diagonal()

Scalar method identical to the corresponding array attribute.

Please see ndarray.diagonal.

dtype

Get array data-descriptor.

dump()

Scalar method identical to the corresponding array attribute.

Please see ndarray.dump.

dumps()

Scalar method identical to the corresponding array attribute.

Please see ndarray.dumps.

fill()

Scalar method identical to the corresponding array attribute.

Please see ndarray.fill.

flags

The integer value of flags.

flat

A 1-D view of the scalar.

flatten()

Scalar method identical to the corresponding array attribute.

Please see ndarray.flatten.

getfield()

Scalar method identical to the corresponding array attribute.

Please see ndarray.getfield.

imag

The imaginary part of the scalar.

item()

Scalar method identical to the corresponding array attribute.

Please see ndarray.item.

itemset
itemsize

The length of one element in bytes.

max()

Scalar method identical to the corresponding array attribute.

Please see ndarray.max.

mean()

Scalar method identical to the corresponding array attribute.

Please see ndarray.mean.

min()

Scalar method identical to the corresponding array attribute.

Please see ndarray.min.

nbytes
ndim

The number of array dimensions.

newbyteorder
nonzero()

Scalar method identical to the corresponding array attribute.

Please see ndarray.nonzero.

prod()

Scalar method identical to the corresponding array attribute.

Please see ndarray.prod.

ptp
put()

Scalar method identical to the corresponding array attribute.

Please see ndarray.put.

ravel()

Scalar method identical to the corresponding array attribute.

Please see ndarray.ravel.

real

The real part of the scalar.

repeat()

Scalar method identical to the corresponding array attribute.

Please see ndarray.repeat.

reshape()

Scalar method identical to the corresponding array attribute.

Please see ndarray.reshape.

resize()

Scalar method identical to the corresponding array attribute.

Please see ndarray.resize.

round()

Scalar method identical to the corresponding array attribute.

Please see ndarray.round.

searchsorted()

Scalar method identical to the corresponding array attribute.

Please see ndarray.searchsorted.

setfield()

Scalar method identical to the corresponding array attribute.

Please see ndarray.setfield.

setflags()

Scalar method identical to the corresponding array attribute.

Please see ndarray.setflags.

shape

Tuple of array dimensions.

size

The number of elements in the gentype.

sort()

Scalar method identical to the corresponding array attribute.

Please see ndarray.sort.

squeeze()

Scalar method identical to the corresponding array attribute.

Please see ndarray.squeeze.

std()

Scalar method identical to the corresponding array attribute.

Please see ndarray.std.

strides

Tuple of bytes steps in each dimension.

sum()

Scalar method identical to the corresponding array attribute.

Please see ndarray.sum.

swapaxes()

Scalar method identical to the corresponding array attribute.

Please see ndarray.swapaxes.

take()

Scalar method identical to the corresponding array attribute.

Please see ndarray.take.

to_device()
tobytes()
tofile()

Scalar method identical to the corresponding array attribute.

Please see ndarray.tofile.

tolist()

Scalar method identical to the corresponding array attribute.

Please see ndarray.tolist.

tostring()

Scalar method identical to the corresponding array attribute.

Please see ndarray.tostring.

trace()

Scalar method identical to the corresponding array attribute.

Please see ndarray.trace.

transpose()

Scalar method identical to the corresponding array attribute.

Please see ndarray.transpose.

var()

Scalar method identical to the corresponding array attribute.

Please see ndarray.var.

view()

Scalar method identical to the corresponding array attribute.

Please see ndarray.view.

quchip.declarative.qnp.geomspace(start, stop, num=50, endpoint=True, dtype=None, axis=0)

Generate geometrically-spaced values.

JAX implementation of numpy.geomspace().

Args:

start: scalar or array. Specifies the starting values. stop: scalar or array. Specifies the stop values. num: int, optional, default=50. Number of values to generate. endpoint: bool, optional, default=True. If True, then include the stop value

in the result. If False, then exclude the stop value.

dtype: optional. Specifies the dtype of the output. axis: int, optional, default=0. Axis along which to generate the geomspace.

Returns:

An array containing the geometrically-spaced values.

See also:

  • jax.numpy.arange(): Generate N evenly-spaced values given a starting point and a step value.

  • jax.numpy.linspace(): Generate evenly-spaced values.

  • jax.numpy.logspace(): Generate logarithmically-spaced values.

Examples:

List 5 geometrically-spaced values between 1 and 16:

>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.geomspace(1, 16, 5)
Array([ 1.,  2.,  4.,  8., 16.], dtype=float32)

List 4 geomtrically-spaced values between 1 and 16, with endpoint=False:

>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.geomspace(1, 16, 4, endpoint=False)
Array([1., 2., 4., 8.], dtype=float32)

Multi-dimensional geomspace:

>>> start = jnp.array([1, 1000])
>>> stop = jnp.array([27, 1])
>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.geomspace(start, stop, 4)
Array([[   1., 1000.],
       [   3.,  100.],
       [   9.,   10.],
       [  27.,    1.]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.get_printoptions()[source]

Alias of numpy.get_printoptions().

JAX arrays are printed via NumPy, so NumPy’s printoptions configurations will apply to printed JAX arrays.

See the numpy.set_printoptions() documentation for details on the available options and their meanings.

quchip.declarative.qnp.gradient(f, *varargs, axis=None, edge_order=None)

Compute the numerical gradient of a sampled function.

JAX implementation of numpy.gradient().

The gradient in jnp.gradient is computed using second-order finite differences across the array of sampled function values. This should not be confused with jax.grad(), which computes a precise gradient of a callable function via automatic differentiation.

Args:

f: N-dimensional array of function values. varargs: optional list of scalars or arrays specifying spacing of

function evaluations. Options are:

  • not specified: unit spacing in all dimensions.

  • a single scalar: constant spacing in all dimensions.

  • N values: specify different spacing in each dimension:

    • scalar values indicate constant spacing in that dimension.

    • array values must match the length of the corresponding dimension, and specify the coordinates at which f is evaluated.

edge_order: not implemented in JAX axis: integer or tuple of integers specifying the axis along which

to compute the gradient. If None (default) calculates the gradient along all axes.

Returns:

an array or tuple of arrays containing the numerical gradient along each specified axis.

See also:

  • jax.grad(): automatic differentiation of a function with a single output.

Examples:

Comparing numerical and automatic differentiation of a simple function:

>>> def f(x):
...   return jnp.sin(x) * jnp.exp(-x / 4)
...
>>> def gradf_exact(x):
...   # exact analytical gradient of f(x)
...   return -f(x) / 4 + jnp.cos(x) * jnp.exp(-x / 4)
...
>>> x = jnp.linspace(0, 5, 10)

>>> with jnp.printoptions(precision=2, suppress=True):
...   print("numerical gradient:", jnp.gradient(f(x), x))
...   print("automatic gradient:", jax.vmap(jax.grad(f))(x))
...   print("exact gradient:    ", gradf_exact(x))
...
numerical gradient: [ 0.83  0.61  0.18 -0.2  -0.43 -0.49 -0.39 -0.21 -0.02  0.08]
automatic gradient: [ 1.    0.62  0.17 -0.23 -0.46 -0.51 -0.41 -0.21 -0.01  0.15]
exact gradient:     [ 1.    0.62  0.17 -0.23 -0.46 -0.51 -0.41 -0.21 -0.01  0.15]

Notice that, as expected, the numerical gradient has some approximation error compared to the automatic gradient computed via jax.grad().

Parameters:
Return type:

Array | list[Array]

quchip.declarative.qnp.greater(x, y, /)

Return element-wise truth value of x > y.

JAX implementation of numpy.greater.

Args:

x: input array or scalar. y: input array or scalar. x and y must either have same shape or be

broadcast compatible.

Returns:

An array containing boolean values. True if the elements of x > y, and False otherwise.

See also:

  • jax.numpy.less(): Returns element-wise truth value of x < y.

  • jax.numpy.greater_equal(): Returns element-wise truth value of x >= y.

  • jax.numpy.less_equal(): Returns element-wise truth value of x <= y.

Examples:

Scalar inputs:

>>> jnp.greater(5, 2)
Array(True, dtype=bool, weak_type=True)

Inputs with same shape:

>>> x = jnp.array([5, 9, -2])
>>> y = jnp.array([4, -1, 6])
>>> jnp.greater(x, y)
Array([ True,  True, False], dtype=bool)

Inputs with broadcast compatibility:

>>> x1 = jnp.array([[5, -6, 7],
...                 [-2, 5, 9]])
>>> y1 = jnp.array([-4, 3, 10])
>>> jnp.greater(x1, y1)
Array([[ True, False, False],
       [ True,  True, False]], dtype=bool)
Parameters:
Return type:

Array

quchip.declarative.qnp.greater_equal(x, y, /)

Return element-wise truth value of x >= y.

JAX implementation of numpy.greater_equal.

Args:

x: input array or scalar. y: input array or scalar. x and y must either have same shape or be

broadcast compatible.

Returns:

An array containing boolean values. True if the elements of x >= y, and False otherwise.

See also:

  • jax.numpy.less_equal(): Returns element-wise truth value of x <= y.

  • jax.numpy.greater(): Returns element-wise truth value of x > y.

  • jax.numpy.less(): Returns element-wise truth value of x < y.

Examples:

Scalar inputs:

>>> jnp.greater_equal(4, 7)
Array(False, dtype=bool, weak_type=True)

Inputs with same shape:

>>> x = jnp.array([2, 5, -1])
>>> y = jnp.array([-6, 4, 3])
>>> jnp.greater_equal(x, y)
Array([ True,  True, False], dtype=bool)

Inputs with broadcast compatibility:

>>> x1 = jnp.array([[3, -1, 4],
...                 [5, 9, -6]])
>>> y1 = jnp.array([-1, 4, 2])
>>> jnp.greater_equal(x1, y1)
Array([[ True, False,  True],
       [ True,  True, False]], dtype=bool)
Parameters:
Return type:

Array

quchip.declarative.qnp.hamming(M)

Return a Hamming window of size M.

JAX implementation of numpy.hamming().

Args:

M: The window size.

Returns:

An array of size M containing the Hamming window.

Examples:
>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jnp.hamming(4))
[0.08 0.77 0.77 0.08]
See also:

  • jax.numpy.bartlett(): return a Bartlett window of size M.

  • jax.numpy.blackman(): return a Blackman window of size M.

  • jax.numpy.hanning(): return a Hanning window of size M.

  • jax.numpy.kaiser(): return a Kaiser window of size M.

Parameters:

M (int)

Return type:

Array

quchip.declarative.qnp.hanning(M)

Return a Hanning window of size M.

JAX implementation of numpy.hanning().

Args:

M: The window size.

Returns:

An array of size M containing the Hanning window.

Examples:
>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jnp.hanning(4))
[0.   0.75 0.75 0.  ]
See also:

  • jax.numpy.bartlett(): return a Bartlett window of size M.

  • jax.numpy.blackman(): return a Blackman window of size M.

  • jax.numpy.hamming(): return a Hamming window of size M.

  • jax.numpy.kaiser(): return a Kaiser window of size M.

Parameters:

M (int)

Return type:

Array

quchip.declarative.qnp.heaviside(x1, x2, /)

Compute the heaviside step function.

JAX implementation of numpy.heaviside.

The heaviside step function is defined by:

\[\begin{split}\mathrm{heaviside}(x1, x2) = \begin{cases} 0, & x1 < 0\\ x2, & x1 = 0\\ 1, & x1 > 0. \end{cases}\end{split}\]
Args:

x1: input array or scalar. complex dtype are not supported. x2: scalar or array. Specifies the return values when x1 is 0. complex

dtype are not supported. x1 and x2 must either have same shape or broadcast compatible.

Returns:

An array containing the heaviside step function of x1, promoting to inexact dtype.

Examples:
>>> x1 = jnp.array([[-2, 0, 3],
...                 [5, -1, 0],
...                 [0, 7, -3]])
>>> x2 = jnp.array([2, 0.5, 1])
>>> jnp.heaviside(x1, x2)
Array([[0. , 0.5, 1. ],
       [1. , 0. , 1. ],
       [2. , 1. , 0. ]], dtype=float32)
>>> jnp.heaviside(x1, 0.5)
Array([[0. , 0.5, 1. ],
       [1. , 0. , 0.5],
       [0.5, 1. , 0. ]], dtype=float32)
>>> jnp.heaviside(-3, x2)
Array([0., 0., 0.], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.histogram(a, bins=10, range=None, weights=None, density=None)

Compute a 1-dimensional histogram.

JAX implementation of numpy.histogram().

Args:

a: array of values to be binned. May be any size or dimension. bins: Specify the number of bins in the histogram (default: 10). bins

may also be an array specifying the locations of the bin edges.

range: tuple of scalars. Specifies the range of the data. If not specified,

the range is inferred from the data.

weights: An optional array specifying the weights of the data points.

Should be broadcast-compatible with a. If not specified, each data point is weighted equally.

density: If True, return the normalized histogram in units of counts

per unit length. If False (default) return the (weighted) counts per bin.

Returns:

A tuple of arrays (histogram, bin_edges), where histogram contains the aggregated data, and bin_edges specifies the boundaries of the bins.

See Also:

  • jax.numpy.bincount(): Count the number of occurrences of each value in an array.

  • jax.numpy.histogram2d(): Compute the histogram of a 2D array.

  • jax.numpy.histogramdd(): Compute the histogram of an N-dimensional array.

  • jax.numpy.histogram_bin_edges(): Compute the bin edges for a histogram.

Examples:
>>> a = jnp.array([1, 2, 3, 10, 11, 15, 19, 25])
>>> counts, bin_edges = jnp.histogram(a, bins=8)
>>> print(counts)
[3. 0. 0. 2. 1. 0. 1. 1.]
>>> print(bin_edges)
[ 1.  4.  7. 10. 13. 16. 19. 22. 25.]

Specifying the bin range:

>>> counts, bin_edges = jnp.histogram(a, range=(0, 25), bins=5)
>>> print(counts)
[3. 0. 2. 2. 1.]
>>> print(bin_edges)
[ 0.  5. 10. 15. 20. 25.]

Specifying the bin edges explicitly:

>>> bin_edges = jnp.array([0, 10, 20, 30])
>>> counts, _ = jnp.histogram(a, bins=bin_edges)
>>> print(counts)
[3. 4. 1.]

Using density=True returns a normalized histogram:

>>> density, bin_edges = jnp.histogram(a, density=True)
>>> dx = jnp.diff(bin_edges)
>>> normed_sum = jnp.sum(density * dx)
>>> jnp.allclose(normed_sum, 1.0)
Array(True, dtype=bool)
Parameters:
Return type:

tuple[Array, Array]

quchip.declarative.qnp.histogram2d(x, y, bins=10, range=None, weights=None, density=None)

Compute a 2-dimensional histogram.

JAX implementation of numpy.histogram2d().

Args:

x: one-dimensional array of x-values for points to be binned. y: one-dimensional array of y-values for points to be binned. bins: Specify the number of bins in the histogram (default: 10). bins

may also be an array specifying the locations of the bin edges, or a pair of integers or pair of arrays specifying the number of bins in each dimension.

range: Pair of arrays or lists of the form [[xmin, xmax], [ymin, ymax]]

specifying the range of the data in each dimension. If not specified, the range is inferred from the data.

weights: An optional array specifying the weights of the data points.

Should be the same shape as x and y. If not specified, each data point is weighted equally.

density: If True, return the normalized histogram in units of counts

per unit area. If False (default) return the (weighted) counts per bin.

Returns:

A tuple of arrays (histogram, x_edges, y_edges), where histogram contains the aggregated data, and x_edges and y_edges specify the boundaries of the bins.

See Also:

  • jax.numpy.histogram(): Compute the histogram of a 1D array.

  • jax.numpy.histogramdd(): Compute the histogram of an N-dimensional array.

  • jax.numpy.histogram_bin_edges(): Compute the bin edges for a histogram.

Examples:
>>> x = jnp.array([1, 2, 3, 10, 11, 15, 19, 25])
>>> y = jnp.array([2, 5, 6, 8, 13, 16, 17, 18])
>>> counts, x_edges, y_edges = jnp.histogram2d(x, y, bins=8)
>>> counts.shape
(8, 8)
>>> x_edges
Array([ 1.,  4.,  7., 10., 13., 16., 19., 22., 25.], dtype=float32)
>>> y_edges
Array([ 2.,  4.,  6.,  8., 10., 12., 14., 16., 18.], dtype=float32)

Specifying the bin range:

>>> counts, x_edges, y_edges = jnp.histogram2d(x, y, range=[(0, 25), (0, 25)], bins=5)
>>> counts.shape
(5, 5)
>>> x_edges
Array([ 0.,  5., 10., 15., 20., 25.], dtype=float32)
>>> y_edges
Array([ 0.,  5., 10., 15., 20., 25.], dtype=float32)

Specifying the bin edges explicitly:

>>> x_edges = jnp.array([0, 10, 20, 30])
>>> y_edges = jnp.array([0, 10, 20, 30])
>>> counts, _, _ = jnp.histogram2d(x, y, bins=[x_edges, y_edges])
>>> counts
Array([[3, 0, 0],
       [1, 3, 0],
       [0, 1, 0]], dtype=int32)

Using density=True returns a normalized histogram:

>>> density, x_edges, y_edges = jnp.histogram2d(x, y, density=True)
>>> dx = jnp.diff(x_edges)
>>> dy = jnp.diff(y_edges)
>>> normed_sum = jnp.sum(density * dx[:, None] * dy[None, :])
>>> jnp.allclose(normed_sum, 1.0)
Array(True, dtype=bool)
Parameters:
Return type:

tuple[Array, Array, Array]

quchip.declarative.qnp.histogram_bin_edges(a, bins=10, range=None, weights=None)

Compute the bin edges for a histogram.

JAX implementation of numpy.histogram_bin_edges().

Args:

a: array of values to be binned bins: Specify the number of bins in the histogram (default: 10). range: tuple of scalars. Specifies the range of the data. If not specified,

the range is inferred from the data.

weights: unused by JAX.

Returns:

An array of bin edges for the histogram.

See also:

  • jax.numpy.histogram(): compute a 1D histogram.

  • jax.numpy.histogram2d(): compute a 2D histogram.

  • jax.numpy.histogramdd(): compute an N-dimensional histogram.

Examples:
>>> a = jnp.array([2, 5, 3, 6, 4, 1])
>>> jnp.histogram_bin_edges(a, bins=5)
Array([1., 2., 3., 4., 5., 6.], dtype=float32)
>>> jnp.histogram_bin_edges(a, bins=5, range=(-10, 10))
Array([-10.,  -6.,  -2.,   2.,   6.,  10.], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.histogramdd(sample, bins=10, range=None, weights=None, density=None)

Compute an N-dimensional histogram.

JAX implementation of numpy.histogramdd().

Args:
sample: input array of shape (N, D) representing N points in

D dimensions.

bins: Specify the number of bins in each dimension of the histogram.

(default: 10). May also be a length-D sequence of integers or arrays of bin edges.

range: Length-D sequence of pairs specifying the range for each dimension.

If not specified, the range is inferred from the data.

weights: An optional shape (N,) array specifying the weights of the

data points. Should be the same shape as sample. If not specified, each data point is weighted equally.

density: If True, return the normalized histogram in units of counts

per unit volume. If False (default) return the (weighted) counts per bin.

Returns:

A tuple of arrays (histogram, bin_edges), where histogram contains the aggregated data, and bin_edges specifies the boundaries of the bins.

See Also:

  • jax.numpy.histogram(): Compute the histogram of a 1D array.

  • jax.numpy.histogram2d(): Compute the histogram of a 2D array.

  • jax.numpy.histogram_bin_edges(): Compute the bin edges for a histogram.

Examples:

A histogram over 100 points in three dimensions

>>> key = jax.random.key(42)
>>> a = jax.random.normal(key, (100, 3))
>>> counts, bin_edges = jnp.histogramdd(a, bins=6,
...                                     range=[(-3, 3), (-3, 3), (-3, 3)])
>>> counts.shape
(6, 6, 6)
>>> bin_edges
[Array([-3., -2., -1.,  0.,  1.,  2.,  3.], dtype=float32),
 Array([-3., -2., -1.,  0.,  1.,  2.,  3.], dtype=float32),
 Array([-3., -2., -1.,  0.,  1.,  2.,  3.], dtype=float32)]

Using density=True returns a normalized histogram:

>>> density, bin_edges = jnp.histogramdd(a, density=True)
>>> bin_widths = map(jnp.diff, bin_edges)
>>> dx, dy, dz = jnp.meshgrid(*bin_widths, indexing='ij')
>>> normed = jnp.sum(density * dx * dy * dz)
>>> jnp.allclose(normed, 1.0)
Array(True, dtype=bool)
Parameters:
Return type:

tuple[Array, list[Array]]

quchip.declarative.qnp.hsplit(ary, indices_or_sections)

Split an array into sub-arrays horizontally.

JAX implementation of numpy.hsplit().

Refer to the documentation of jax.numpy.split() for details. hsplit is equivalent to split with axis=1, or axis=0 for one-dimensional arrays.

Examples:

1D array:

>>> x = jnp.array([1, 2, 3, 4, 5, 6])
>>> x1, x2 = jnp.hsplit(x, 2)
>>> print(x1, x2)
[1 2 3] [4 5 6]

2D array:

>>> x = jnp.array([[1, 2, 3, 4],
...                [5, 6, 7, 8]])
>>> x1, x2 = jnp.hsplit(x, 2)
>>> print(x1)
[[1 2]
 [5 6]]
>>> print(x2)
[[3 4]
 [7 8]]
See also:

  • jax.numpy.split(): split an array along any axis.

  • jax.numpy.vsplit(): split vertically, i.e. along axis=0

  • jax.numpy.dsplit(): split depth-wise, i.e. along axis=2

  • jax.numpy.array_split(): like split, but allows indices_or_sections to be an integer that does not evenly divide the size of the array.

Parameters:
Return type:

list[Array]

quchip.declarative.qnp.hstack(tup, dtype=None)

Horizontally stack arrays.

JAX implementation of numpy.hstack().

For arrays of one or more dimensions, this is equivalent to jax.numpy.concatenate() with axis=1.

Args:
tup: a sequence of arrays to stack; each must have the same shape along all

but the second axis. Input arrays will be promoted to at least rank 1. If a single array is given it will be treated equivalently to tup = unstack(tup), but the implementation will avoid explicit unstacking.

dtype: optional dtype of the resulting array. If not specified, the dtype

will be determined via type promotion rules described in type-promotion.

Returns:

the stacked result.

See also:

  • jax.numpy.stack(): stack along arbitrary axes

  • jax.numpy.concatenate(): concatenation along existing axes.

  • jax.numpy.vstack(): stack vertically, i.e. along axis 0.

  • jax.numpy.dstack(): stack depth-wise, i.e. along axis 2.

Examples:

Scalar values:

>>> jnp.hstack([1, 2, 3])
Array([1, 2, 3], dtype=int32, weak_type=True)

1D arrays:

>>> x = jnp.arange(3)
>>> y = jnp.ones(3)
>>> jnp.hstack([x, y])
Array([0., 1., 2., 1., 1., 1.], dtype=float32)

2D arrays:

>>> x = x.reshape(3, 1)
>>> y = y.reshape(3, 1)
>>> jnp.hstack([x, y])
Array([[0., 1.],
       [1., 1.],
       [2., 1.]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.hypot(x1, x2, /)

Return element-wise hypotenuse for the given legs of a right angle triangle.

JAX implementation of numpy.hypot.

Args:
x1: scalar or array. Specifies one of the legs of right angle triangle.

complex dtype are not supported.

x2: scalar or array. Specifies the other leg of right angle triangle.

complex dtype are not supported. x1 and x2 must either have same shape or be broadcast compatible.

Returns:

An array containing the hypotenuse for the given given legs x1 and x2 of a right angle triangle, promoting to inexact dtype.

Note:

jnp.hypot is a more numerically stable way of computing jnp.sqrt(x1 ** 2 + x2 **2).

Examples:
>>> jnp.hypot(3, 4)
Array(5., dtype=float32, weak_type=True)
>>> x1 = jnp.array([[3, -2, 5],
...                 [9, 1, -4]])
>>> x2 = jnp.array([-5, 6, 8])
>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.hypot(x1, x2)
Array([[ 5.831,  6.325,  9.434],
       [10.296,  6.083,  8.944]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.i0(x)

Calculate modified Bessel function of first kind, zeroth order.

JAX implementation of numpy.i0().

Modified Bessel function of first kind, zeroth order is defined by:

\[\mathrm{i0}(x) = I_0(x) = \sum_{k=0}^{\infty} \frac{(x^2/4)^k}{(k!)^2}\]
Args:
x: scalar or array. Specifies the argument of Bessel function. Complex inputs

are not supported.

Returns:

An array containing the corresponding values of the modified Bessel function of x.

See also:

  • jax.scipy.special.i0(): Calculates the modified Bessel function of zeroth order.

  • jax.scipy.special.i1(): Calculates the modified Bessel function of first order.

  • jax.scipy.special.i0e(): Calculates the exponentially scaled modified Bessel function of zeroth order.

Examples:
>>> x = jnp.array([-2, -1, 0, 1, 2])
>>> jnp.i0(x)
Array([2.2795851, 1.266066 , 1.0000001, 1.266066 , 2.2795851], dtype=float32)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.identity(n, dtype=None)

Create a square identity matrix

JAX implementation of numpy.identity().

Args:

n: integer specifying the size of each array dimension. dtype: optional dtype; defaults to floating point.

Returns:

Identity array of shape (n, n).

See also:

jax.numpy.eye(): non-square and/or offset identity matrices.

Examples:

A simple 3x3 identity matrix:

>>> jnp.identity(3)
Array([[1., 0., 0.],
       [0., 1., 0.],
       [0., 0., 1.]], dtype=float32)

A 2x2 integer identity matrix:

>>> jnp.identity(2, dtype=int)
Array([[1, 0],
       [0, 1]], dtype=int32)
Parameters:
Return type:

Array

class quchip.declarative.qnp.iinfo(int_type)[source]

Bases: object

dtype: dtype
kind: str
bits: int
min: int
max: int
quchip.declarative.qnp.imag(val, /)

Return element-wise imaginary of part of the complex argument.

JAX implementation of numpy.imag.

Args:

val: input array or scalar.

Returns:

An array containing the imaginary part of the elements of val.

See also:

  • jax.numpy.conjugate() and jax.numpy.conj(): Returns the element-wise complex-conjugate of the input.

  • jax.numpy.real(): Returns the element-wise real part of the complex argument.

Examples:
>>> jnp.imag(4)
Array(0, dtype=int32, weak_type=True)
>>> jnp.imag(5j)
Array(5., dtype=float32, weak_type=True)
>>> x = jnp.array([2+3j, 5-1j, -3])
>>> jnp.imag(x)
Array([ 3., -1.,  0.], dtype=float32)
Parameters:

val (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.indices(dimensions, dtype=None, sparse=False)

Generate arrays of grid indices.

JAX implementation of numpy.indices().

Args:

dimensions: the shape of the grid. dtype: the dtype of the indices (defaults to integer). sparse: if True, then return sparse indices. Default is False, which

returns dense indices.

Returns:

An array of shape (len(dimensions), *dimensions) If sparse is False, or a sequence of arrays of the same length as dimensions if sparse is True.

See also:

  • jax.numpy.meshgrid(): generate a grid from arbitrary input arrays.

  • jax.numpy.mgrid: generate dense indices using a slicing syntax.

  • jax.numpy.ogrid: generate sparse indices using a slicing syntax.

Examples:
>>> jnp.indices((2, 3))
Array([[[0, 0, 0],
        [1, 1, 1]],

       [[0, 1, 2],
        [0, 1, 2]]], dtype=int32)
>>> jnp.indices((2, 3), sparse=True)
(Array([[0],
       [1]], dtype=int32), Array([[0, 1, 2]], dtype=int32))
Parameters:
Return type:

Array | tuple[Array, …]

class quchip.declarative.qnp.inexact

Bases: number

Abstract base class of all numeric scalar types with a (potentially) inexact representation of the values in its range, such as floating-point numbers.

quchip.declarative.qnp.inner(a, b, *, precision=None, preferred_element_type=None)

Compute the inner product of two arrays.

JAX implementation of numpy.inner().

Unlike jax.numpy.matmul() or jax.numpy.dot(), this always performs a contraction along the last dimension of each input.

Args:

a: array of shape (..., N) b: array of shape (..., N) precision: either None (default), which means the default precision for

the backend, a Precision enum value (Precision.DEFAULT, Precision.HIGH or Precision.HIGHEST) or a tuple of two such values indicating precision of a and b.

preferred_element_type: either None (default), which means the default

accumulation type for the input types, or a datatype, indicating to accumulate results to and return a result with that datatype.

Returns:

array of shape (*a.shape[:-1], *b.shape[:-1]) containing the batched vector product of the inputs.

See also:

  • jax.numpy.vecdot(): conjugate multiplication along a specified axis.

  • jax.numpy.tensordot(): general tensor multiplication.

  • jax.numpy.matmul(): general batched matrix & vector multiplication.

Examples:

For 1D inputs, this implements standard (non-conjugate) vector multiplication:

>>> a = jnp.array([1j, 3j, 4j])
>>> b = jnp.array([4., 2., 5.])
>>> jnp.inner(a, b)
Array(0.+30.j, dtype=complex64)

For multi-dimensional inputs, batch dimensions are stacked rather than broadcast:

>>> a = jnp.ones((2, 3))
>>> b = jnp.ones((5, 3))
>>> jnp.inner(a, b).shape
(2, 5)
Parameters:
Return type:

Array

quchip.declarative.qnp.insert(arr, obj, values, axis=None)

Insert entries into an array at specified indices.

JAX implementation of numpy.insert().

Args:

arr: array object into which values will be inserted. obj: slice or array of indices specifying insertion locations. values: array of values to be inserted. axis: specify the insertion axis in the case of multi-dimensional

arrays. If unspecified, arr will be flattened.

Returns:

A copy of arr with values inserted at the specified locations.

See also:

  • jax.numpy.delete(): delete entries from an array.

Examples:

Inserting a single value:

>>> x = jnp.arange(5)
>>> jnp.insert(x, 2, 99)
Array([ 0,  1, 99,  2,  3,  4], dtype=int32)

Inserting multiple identical values using a slice:

>>> jnp.insert(x, slice(None, None, 2), -1)
Array([-1,  0,  1, -1,  2,  3, -1,  4], dtype=int32)

Inserting multiple values using an index:

>>> indices = jnp.array([4, 2, 5])
>>> values = jnp.array([10, 11, 12])
>>> jnp.insert(x, indices, values)
Array([ 0,  1, 11,  2,  3, 10,  4, 12], dtype=int32)

Inserting columns into a 2D array:

>>> x = jnp.array([[1, 2, 3],
...                [4, 5, 6]])
>>> indices = jnp.array([1, 3])
>>> values = jnp.array([[10, 11],
...                     [12, 13]])
>>> jnp.insert(x, indices, values, axis=1)
Array([[ 1, 10,  2,  3, 11],
       [ 4, 12,  5,  6, 13]], dtype=int32)
Parameters:
Return type:

Array

class quchip.declarative.qnp.int16(x)

Bases: object

A JAX scalar constructor of type int16.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype('int16')
class quchip.declarative.qnp.int2(x)

Bases: object

A JAX scalar constructor of type int2.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype(int2)
class quchip.declarative.qnp.int32(x)

Bases: object

A JAX scalar constructor of type int32.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype('int32')
class quchip.declarative.qnp.int4(x)

Bases: object

A JAX scalar constructor of type int4.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype(int4)
class quchip.declarative.qnp.int64(x)

Bases: object

A JAX scalar constructor of type int64.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype('int64')
class quchip.declarative.qnp.int8(x)

Bases: object

A JAX scalar constructor of type int8.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype('int8')
quchip.declarative.qnp.int_

alias of int64

class quchip.declarative.qnp.integer

Bases: number

Abstract base class of all integer scalar types.

denominator

denominator of value (1)

is_integer() → bool

Return True if the number is finite with integral value.

Added in version 1.22.

Examples

>>> import numpy as np
>>> np.int64(-2).is_integer()
True
>>> np.uint32(5).is_integer()
True
numerator

numerator of value (the value itself)

quchip.declarative.qnp.interp(x, xp, fp, left=None, right=None, period=None)

One-dimensional linear interpolation.

JAX implementation of numpy.interp().

Args:

x: N-dimensional array of x coordinates at which to evaluate the interpolation. xp: one-dimensional sorted array of points to be interpolated. fp: array of shape xp.shape containing the function values associated with xp. left: specify how to handle points x < xp[0]. Default is to return fp[0].

If left is a scalar value, it will return this value. if left is the string "extrapolate", then the value will be determined by linear extrapolation. left is ignored if period is specified.

right: specify how to handle points x > xp[-1]. Default is to return fp[-1].

If right is a scalar value, it will return this value. if right is the string "extrapolate", then the value will be determined by linear extrapolation. right is ignored if period is specified.

period: optionally specify the period for the x coordinates, for e.g. interpolation

in angular space.

Returns:

an array of shape x.shape containing the interpolated function at values x.

Examples:
>>> xp = jnp.arange(10)
>>> fp = 2 * xp
>>> x = jnp.array([0.5, 2.0, 3.5])
>>> interp(x, xp, fp)
Array([1., 4., 7.], dtype=float32)

Unless otherwise specified, extrapolation will be constant:

>>> x = jnp.array([-10., 10.])
>>> interp(x, xp, fp)
Array([ 0., 18.], dtype=float32)

Use "extrapolate" mode for linear extrapolation:

>>> interp(x, xp, fp, left='extrapolate', right='extrapolate')
Array([-20.,  20.], dtype=float32)

For periodic interpolation, specify the period:

>>> xp = jnp.array([0, jnp.pi / 2, jnp.pi, 3 * jnp.pi / 2])
>>> fp = jnp.sin(xp)
>>> x = 2 * jnp.pi  # note: not in input array
>>> jnp.interp(x, xp, fp, period=2 * jnp.pi)
Array(0., dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.intersect1d(ar1, ar2, assume_unique=False, return_indices=False, *, size=None, fill_value=None)

Compute the set intersection of two 1D arrays.

JAX implementation of numpy.intersect1d().

Because the size of the output of intersect1d is data-dependent, the function is not typically compatible with jit() and other JAX transformations. The JAX version adds the optional size argument which must be specified statically for jnp.intersect1d to be used in such contexts.

Args:

ar1: first array of values to intersect. ar2: second array of values to intersect. assume_unique: if True, assume the input arrays contain unique values. This allows

a more efficient implementation, but if assume_unique is True and the input arrays contain duplicates, the behavior is undefined. default: False.

return_indices: If True, return arrays of indices specifying where the intersected

values first appear in the input arrays.

size: if specified, return only the first size sorted elements. If there are fewer

elements than size indicates, the return value will be padded with fill_value, and returned indices will be padded with an out-of-bound index.

fill_value: when size is specified and there are fewer than the indicated number of

elements, fill the remaining entries fill_value. Defaults to the smallest value in the intersection.

Returns:

An array intersection, or if return_indices=True, a tuple of arrays (intersection, ar1_indices, ar2_indices). Returned values are

  • intersection: A 1D array containing each value that appears in both ar1 and ar2.

  • ar1_indices: (returned if return_indices=True) an array of shape intersection.shape containing the indices in flattened ar1 of values in intersection. For 1D inputs, intersection is equivalent to ar1[ar1_indices].

  • ar2_indices: (returned if return_indices=True) an array of shape intersection.shape containing the indices in flattened ar2 of values in intersection. For 1D inputs, intersection is equivalent to ar2[ar2_indices].

See also:

  • jax.numpy.union1d(): the set union of two 1D arrays.

  • jax.numpy.setxor1d(): the set XOR of two 1D arrays.

  • jax.numpy.setdiff1d(): the set difference of two 1D arrays.

Examples:
>>> ar1 = jnp.array([1, 2, 3, 4])
>>> ar2 = jnp.array([3, 4, 5, 6])
>>> jnp.intersect1d(ar1, ar2)
Array([3, 4], dtype=int32)

Computing intersection with indices:

>>> intersection, ar1_indices, ar2_indices = jnp.intersect1d(ar1, ar2, return_indices=True)
>>> intersection
Array([3, 4], dtype=int32)

ar1_indices gives the indices of the intersected values within ar1:

>>> ar1_indices
Array([2, 3], dtype=int32)
>>> jnp.all(intersection == ar1[ar1_indices])
Array(True, dtype=bool)

ar2_indices gives the indices of the intersected values within ar2:

>>> ar2_indices
Array([0, 1], dtype=int32)
>>> jnp.all(intersection == ar2[ar2_indices])
Array(True, dtype=bool)
Parameters:
Return type:

Array | tuple[Array, Array, Array]

quchip.declarative.qnp.invert(x, /)

Compute the bitwise inversion of an input.

JAX implementation of numpy.invert(). This function provides the implementation of the ~ operator for JAX arrays.

Args:

x: input array, must be boolean or integer typed.

Returns:

An array of the same shape and dtype as `x, with the bits inverted.

See also:

  • jax.numpy.bitwise_invert(): Array API alias of this function.

  • jax.numpy.logical_not(): Invert after casting input to boolean.

Examples:
>>> x = jnp.arange(5, dtype='uint8')
>>> print(x)
[0 1 2 3 4]
>>> print(jnp.invert(x))
[255 254 253 252 251]

This function implements the unary ~ operator for JAX arrays:

>>> print(~x)
[255 254 253 252 251]

invert() operates bitwise on the input, and so the meaning of its output may be more clear by showing the bitwise representation:

>>> with jnp.printoptions(formatter={'int': lambda x: format(x, '#010b')}):
...   print(f"{x  = }")
...   print(f"{~x = }")
x  = Array([0b00000000, 0b00000001, 0b00000010, 0b00000011, 0b00000100], dtype=uint8)
~x = Array([0b11111111, 0b11111110, 0b11111101, 0b11111100, 0b11111011], dtype=uint8)

For boolean inputs, invert() is equivalent to logical_not():

>>> x = jnp.array([True, False, True, True, False])
>>> jnp.invert(x)
Array([False,  True, False, False,  True], dtype=bool)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.isclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False)

Check if the elements of two arrays are approximately equal within a tolerance.

JAX implementation of numpy.allclose().

Essentially this function evaluates the following condition:

\[|a - b| \le \mathtt{atol} + \mathtt{rtol} * |b|\]

jnp.inf in a will be considered equal to jnp.inf in b.

Args:

a: first input array to compare. b: second input array to compare. rtol: relative tolerance used for approximate equality. Default = 1e-05. atol: absolute tolerance used for approximate equality. Default = 1e-08. equal_nan: Boolean. If True, NaNs in a will be considered

equal to NaNs in b. Default is False.

Returns:

A new array containing boolean values indicating whether the input arrays are element-wise approximately equal within the specified tolerances.

See Also:

  • jax.numpy.allclose()

  • jax.numpy.equal()

Examples:
>>> jnp.isclose(jnp.array([1e6, 2e6, jnp.inf]), jnp.array([1e6, 2e7, jnp.inf]))
Array([ True, False,  True], dtype=bool)
>>> jnp.isclose(jnp.array([1e6, 2e6, 3e6]),
...              jnp.array([1.00008e6, 2.00008e7, 3.00008e8]), rtol=1e3)
Array([ True,  True,  True], dtype=bool)
>>> jnp.isclose(jnp.array([1e6, 2e6, 3e6]),
...              jnp.array([1.00001e6, 2.00002e6, 3.00009e6]), atol=1e3)
Array([ True,  True,  True], dtype=bool)
>>> jnp.isclose(jnp.array([jnp.nan, 1, 2]),
...              jnp.array([jnp.nan, 1, 2]), equal_nan=True)
Array([ True,  True,  True], dtype=bool)
Parameters:
Return type:

Array

quchip.declarative.qnp.iscomplex(x)

Return boolean array showing where the input is complex.

JAX implementation of numpy.iscomplex().

Args:

x: Input array to check.

Returns:

A new array containing boolean values indicating complex elements.

See Also:

  • jax.numpy.iscomplexobj()

  • jax.numpy.isrealobj()

Examples:
>>> jnp.iscomplex(jnp.array([True, 0, 1, 2j, 1+2j]))
Array([False, False, False, True, True], dtype=bool)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.iscomplexobj(x)

Check if the input is a complex number or an array containing complex elements.

JAX implementation of numpy.iscomplexobj().

The function evaluates based on input type rather than value. Inputs with zero imaginary parts are still considered complex.

Args:

x: input object to check.

Returns:

True if x is a complex number or an array containing at least one complex element, False otherwise.

See Also:

  • jax.numpy.isrealobj()

  • jax.numpy.iscomplex()

Examples:
>>> jnp.iscomplexobj(True)
False
>>> jnp.iscomplexobj(0)
False
>>> jnp.iscomplexobj(jnp.array([1, 2]))
False
>>> jnp.iscomplexobj(1+2j)
True
>>> jnp.iscomplexobj(jnp.array([0, 1+2j]))
True
Parameters:

x (Any)

Return type:

bool

quchip.declarative.qnp.isdtype(dtype, kind)

Returns a boolean indicating whether a provided dtype is of a specified kind.

Args:

dtype : the input dtype kind : the data type kind.

If kind is dtype-like, return dtype = kind. If kind is a string, then return True if the dtype is in the specified category:

  • 'bool': {bool}

  • 'signed integer': {int4, int8, int16, int32, int64}

  • 'unsigned integer': {uint4, uint8, uint16, uint32, uint64}

  • 'integral': shorthand for ('signed integer', 'unsigned integer')

  • 'real floating': {float8_*, float16, bfloat16, float32, float64}

  • 'complex floating': {complex64, complex128}

  • 'numeric': shorthand for ('integral', 'real floating', 'complex floating')

If kind is a tuple, then return True if dtype matches any entry of the tuple.

Returns:

True or False

Parameters:
Return type:

bool

quchip.declarative.qnp.isfinite(x, /)

Return a boolean array indicating whether each element of input is finite.

JAX implementation of numpy.isfinite.

Args:

x: input array or scalar.

Returns:

A boolean array of same shape as x containing True where x is not inf, -inf, or NaN, and False otherwise.

See also:

  • jax.numpy.isinf(): Returns a boolean array indicating whether each element of input is either positive or negative infinity.

  • jax.numpy.isposinf(): Returns a boolean array indicating whether each element of input is positive infinity.

  • jax.numpy.isneginf(): Returns a boolean array indicating whether each element of input is negative infinity.

  • jax.numpy.isnan(): Returns a boolean array indicating whether each element of input is not a number (NaN).

Examples:
>>> x = jnp.array([-1, 3, jnp.inf, jnp.nan])
>>> jnp.isfinite(x)
Array([ True,  True, False, False], dtype=bool)
>>> jnp.isfinite(3-4j)
Array(True, dtype=bool, weak_type=True)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.isin(element, test_elements, assume_unique=False, invert=False, *, method='auto')

Determine whether elements in element appear in test_elements.

JAX implementation of numpy.isin().

Args:

element: input array of elements for which membership will be checked. test_elements: N-dimensional array of test values to check for the presence of

each element.

invert: If True, return ~isin(element, test_elements). Default is False. assume_unique: if true, input arrays are assumed to be unique, which can

lead to more efficient computation. If the input arrays are not unique and assume_unique is set to True, the results are undefined.

method: string specifying the method used to compute the result. Supported

options are ‘compare_all’, ‘binary_search’, ‘sort’, and ‘auto’ (default).

Returns:

A boolean array of shape element.shape that specifies whether each element appears in test_elements.

Examples:
>>> elements = jnp.array([1, 2, 3, 4])
>>> test_elements = jnp.array([[1, 5, 6, 3, 7, 1]])
>>> jnp.isin(elements, test_elements)
Array([ True, False,  True, False], dtype=bool)
Parameters:
Return type:

Array

quchip.declarative.qnp.isinf(x, /)

Return a boolean array indicating whether each element of input is infinite.

JAX implementation of numpy.isinf.

Args:

x: input array or scalar.

Returns:

A boolean array of same shape as x containing True where x is inf or -inf, and False otherwise.

See also:

  • jax.numpy.isposinf(): Returns a boolean array indicating whether each element of input is positive infinity.

  • jax.numpy.isneginf(): Returns a boolean array indicating whether each element of input is negative infinity.

  • jax.numpy.isfinite(): Returns a boolean array indicating whether each element of input is finite.

  • jax.numpy.isnan(): Returns a boolean array indicating whether each element of input is not a number (NaN).

Examples:
>>> jnp.isinf(jnp.inf)
Array(True, dtype=bool)
>>> x = jnp.array([2+3j, -jnp.inf, 6, jnp.inf, jnp.nan])
>>> jnp.isinf(x)
Array([False,  True, False,  True, False], dtype=bool)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.isnan(x, /)

Returns a boolean array indicating whether each element of input is NaN.

JAX implementation of numpy.isnan.

Args:

x: input array or scalar.

Returns:

A boolean array of same shape as x containing True where x is not a number (i.e. NaN) and False otherwise.

See also:

  • jax.numpy.isfinite(): Returns a boolean array indicating whether each element of input is finite.

  • jax.numpy.isinf(): Returns a boolean array indicating whether each element of input is either positive or negative infinity.

  • jax.numpy.isposinf(): Returns a boolean array indicating whether each element of input is positive infinity.

  • jax.numpy.isneginf(): Returns a boolean array indicating whether each element of input is negative infinity.

Examples:
>>> jnp.isnan(6)
Array(False, dtype=bool, weak_type=True)
>>> x = jnp.array([2, 1+4j, jnp.inf, jnp.nan])
>>> jnp.isnan(x)
Array([False, False, False,  True], dtype=bool)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.isneginf(x, /, out=None)

Return boolean array indicating whether each element of input is negative infinite.

JAX implementation of numpy.isneginf.

Args:

x: input array or scalar. complex dtype are not supported.

Returns:

A boolean array of same shape as x containing True where x is -inf, and False otherwise.

See also:

  • jax.numpy.isinf(): Returns a boolean array indicating whether each element of input is either positive or negative infinity.

  • jax.numpy.isposinf(): Returns a boolean array indicating whether each element of input is positive infinity.

  • jax.numpy.isfinite(): Returns a boolean array indicating whether each element of input is finite.

  • jax.numpy.isnan(): Returns a boolean array indicating whether each element of input is not a number (NaN).

Examples:
>>> jnp.isneginf(jnp.inf)
Array(False, dtype=bool)
>>> x = jnp.array([-jnp.inf, 5, jnp.inf, jnp.nan, 1])
>>> jnp.isneginf(x)
Array([ True, False, False, False, False], dtype=bool)
quchip.declarative.qnp.isposinf(x, /, out=None)

Return boolean array indicating whether each element of input is positive infinite.

JAX implementation of numpy.isposinf.

Args:

x: input array or scalar. complex dtype are not supported.

Returns:

A boolean array of same shape as x containing True where x is inf, and False otherwise.

See also:

  • jax.numpy.isinf(): Returns a boolean array indicating whether each element of input is either positive or negative infinity.

  • jax.numpy.isneginf(): Returns a boolean array indicating whether each element of input is negative infinity.

  • jax.numpy.isfinite(): Returns a boolean array indicating whether each element of input is finite.

  • jax.numpy.isnan(): Returns a boolean array indicating whether each element of input is not a number (NaN).

Examples:
>>> jnp.isposinf(5)
Array(False, dtype=bool)
>>> x = jnp.array([-jnp.inf, 5, jnp.inf, jnp.nan, 1])
>>> jnp.isposinf(x)
Array([False, False,  True, False, False], dtype=bool)
quchip.declarative.qnp.isreal(x)

Return boolean array showing where the input is real.

JAX implementation of numpy.isreal().

Args:

x: input array to check.

Returns:

A new array containing boolean values indicating real elements.

See Also:

  • jax.numpy.iscomplex()

  • jax.numpy.isrealobj()

Examples:
>>> jnp.isreal(jnp.array([False, 0j, 1, 2.1, 1+2j]))
Array([ True,  True,  True,  True, False], dtype=bool)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.isrealobj(x)

Check if the input is not a complex number or an array containing complex elements.

JAX implementation of numpy.isrealobj().

The function evaluates based on input type rather than value. Inputs with zero imaginary parts are still considered complex.

Args:

x: input object to check.

Returns:

False if x is a complex number or an array containing at least one complex element, True otherwise.

See Also:

  • jax.numpy.iscomplexobj()

  • jax.numpy.isreal()

Examples:
>>> jnp.isrealobj(0)
True
>>> jnp.isrealobj(1.2)
True
>>> jnp.isrealobj(jnp.array([1, 2]))
True
>>> jnp.isrealobj(1+2j)
False
>>> jnp.isrealobj(jnp.array([0, 1+2j]))
False
Parameters:

x (Any)

Return type:

bool

quchip.declarative.qnp.isscalar(element)

Return True if the input is a scalar.

JAX implementation of numpy.isscalar(). JAX’s implementation differs from NumPy’s in that it considers zero-dimensional arrays to be scalars; see the Note below for more details.

Args:

element: input object to check; any type is valid input.

Returns:

True if element is a scalar value or an array-like object with zero dimensions, False otherwise.

Note:

JAX and NumPy differ in their representation of scalar values. NumPy has special scalar objects (e.g. np.int32(0)) which are distinct from zero-dimensional arrays (e.g. np.array(0)), and numpy.isscalar() returns True for the former and False for the latter.

JAX does not define special scalar objects, but rather represents scalars as zero-dimensional arrays. As such, jax.numpy.isscalar() returns True for both scalar objects (e.g. 0.0 or np.float32(0.0)) and array-like objects with zero dimensions (e.g. jnp.array(0.0), np.array(0.0)).

One reason for the different conventions in isscalar is to maintain JIT-invariance: i.e. the property that the result of a function should not change when it is JIT-compiled. Because scalar inputs are cast to zero-dimensional JAX arrays at JIT boundaries, the semantics of numpy.isscalar() are such that the result changes under JIT:

>>> np.isscalar(1.0)
True
>>> jax.jit(np.isscalar)(1.0)
Array(False, dtype=bool)

By treating zero-dimensional arrays as scalars, jax.numpy.isscalar() avoids this issue:

>>> jnp.isscalar(1.0)
True
>>> jax.jit(jnp.isscalar)(1.0)
Array(True, dtype=bool)
Examples:

In JAX, both scalars and zero-dimensional array-like objects are considered scalars:

>>> jnp.isscalar(1.0)
True
>>> jnp.isscalar(1 + 1j)
True
>>> jnp.isscalar(jnp.array(1))  # zero-dimensional JAX array
True
>>> jnp.isscalar(jnp.int32(1))  # JAX scalar constructor
True
>>> jnp.isscalar(np.array(1.0))  # zero-dimensional NumPy array
True
>>> jnp.isscalar(np.int32(1))  # NumPy scalar type
True

Arrays with one or more dimension are not considered scalars:

>>> jnp.isscalar(jnp.array([1]))
False
>>> jnp.isscalar(np.array([1]))
False

Compare this to numpy.isscalar(), which returns True for scalar-typed objects, and False for all arrays, even those with zero dimensions:

>>> np.isscalar(np.int32(1))  # scalar object
True
>>> np.isscalar(np.array(1))  # zero-dimensional array
False

In JAX, as in NumPy, objects which are not array-like are not considered scalars:

>>> jnp.isscalar(None)
False
>>> jnp.isscalar([1])
False
>>> jnp.isscalar(tuple())
False
>>> jnp.isscalar(slice(10))
False
Parameters:

element (Any)

Return type:

bool

quchip.declarative.qnp.issubdtype(arg1, arg2)

Return True if arg1 is equal or lower than arg2 in the type hierarchy.

JAX implementation of numpy.issubdtype().

The main difference in JAX’s implementation is that it properly handles dtype extensions such as bfloat16.

Args:
arg1: dtype-like object. In typical usage, this will be a dtype specifier,

such as "float32" (i.e. a string), np.dtype('int32') (i.e. an instance of numpy.dtype), jnp.complex64 (i.e. a JAX scalar constructor), or np.uint8 (i.e. a NumPy scalar type).

arg2: dtype-like object. In typical usage, this will be a generic scalar

type, such as jnp.integer, jnp.floating, or jnp.complexfloating.

Returns:

True if arg1 represents a dtype that is equal or lower in the type hierarchy than arg2.

See also:

  • jax.numpy.isdtype(): similar function aligning with the array API standard.

Examples:
>>> jnp.issubdtype('uint32', jnp.unsignedinteger)
True
>>> jnp.issubdtype(np.int32, jnp.integer)
True
>>> jnp.issubdtype(jnp.bfloat16, jnp.floating)
True
>>> jnp.issubdtype(np.dtype('complex64'), jnp.complexfloating)
True
>>> jnp.issubdtype('complex64', jnp.integer)
False

Be aware that while this is very similar to numpy.issubdtype(), the results of these differ in the case of JAX’s custom floating point types:

>>> np.issubdtype('bfloat16', np.floating)
False
>>> jnp.issubdtype('bfloat16', jnp.floating)
True
Parameters:
Return type:

bool

quchip.declarative.qnp.iterable(y)

Check whether or not an object can be iterated over.

Parameters:

y (object) – Input object.

Returns:

b – Return True if the object has an iterator method or is a sequence and False otherwise.

Return type:

bool

Examples

>>> import numpy as np
>>> np.iterable([1, 2, 3])
True
>>> np.iterable(2)
False

Notes

In most cases, the results of np.iterable(obj) are consistent with isinstance(obj, collections.abc.Iterable). One notable exception is the treatment of 0-dimensional arrays:

>>> from collections.abc import Iterable
>>> a = np.array(1.0)  # 0-dimensional numpy array
>>> isinstance(a, Iterable)
True
>>> np.iterable(a)
False
quchip.declarative.qnp.ix_(*args)

Return a multi-dimensional grid (open mesh) from N one-dimensional sequences.

JAX implementation of numpy.ix_().

Args:

*args: N one-dimensional arrays

Returns:

Tuple of Jax arrays forming an open mesh, each with N dimensions.

See Also:

  • jax.numpy.ogrid

  • jax.numpy.mgrid

  • jax.numpy.meshgrid()

Examples:
>>> rows = jnp.array([0, 2])
>>> cols = jnp.array([1, 3])
>>> open_mesh = jnp.ix_(rows, cols)
>>> open_mesh
(Array([[0],
      [2]], dtype=int32), Array([[1, 3]], dtype=int32))
>>> [grid.shape for grid in open_mesh]
[(2, 1), (1, 2)]
>>> x = jnp.array([[10, 20, 30, 40],
...                [50, 60, 70, 80],
...                [90, 100, 110, 120],
...                [130, 140, 150, 160]])
>>> x[open_mesh]
Array([[ 20,  40],
       [100, 120]], dtype=int32)
Parameters:

args (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

tuple[Array, …]

quchip.declarative.qnp.kaiser(M, beta)

Return a Kaiser window of size M.

JAX implementation of numpy.kaiser().

Args:

M: The window size. beta: The Kaiser window parameter.

Returns:

An array of size M containing the Kaiser window.

Examples:
>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jnp.kaiser(4, 1.5))
[0.61 0.95 0.95 0.61]
See also:

  • jax.numpy.bartlett(): return a Bartlett window of size M.

  • jax.numpy.blackman(): return a Blackman window of size M.

  • jax.numpy.hamming(): return a Hamming window of size M.

  • jax.numpy.hanning(): return a Hanning window of size M.

Parameters:
Return type:

Array

quchip.declarative.qnp.kron(a, b)

Compute the Kronecker product of two input arrays.

JAX implementation of numpy.kron().

The Kronecker product is an operation on two matrices of arbitrary size that produces a block matrix. Each element of the first matrix a is multiplied by the entire second matrix b. If a has shape (m, n) and b has shape (p, q), the resulting matrix will have shape (m * p, n * q).

Args:

a: first input array with any shape. b: second input array with any shape.

Returns:

A new array representing the Kronecker product of the inputs a and b. The shape of the output is the element-wise product of the input shapes.

See also:

  • jax.numpy.outer(): compute the outer product of two arrays.

Examples:
>>> a = jnp.array([[1, 2],
...                [3, 4]])
>>> b = jnp.array([[5, 6],
...                [7, 8]])
>>> jnp.kron(a, b)
Array([[ 5,  6, 10, 12],
       [ 7,  8, 14, 16],
       [15, 18, 20, 24],
       [21, 24, 28, 32]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.lcm(x1, x2)

Compute the least common multiple of two arrays.

JAX implementation of numpy.lcm().

Args:

x1: First input array. The elements must have integer dtype. x2: Second input array. The elements must have integer dtype.

Returns:

An array containing the least common multiple of the corresponding elements from the absolute values of x1 and x2.

See also:

  • jax.numpy.gcd(): compute the greatest common divisor of two arrays.

Examples:

Scalar inputs:

>>> jnp.lcm(12, 18)
Array(36, dtype=int32, weak_type=True)

Array inputs:

>>> x1 = jnp.array([12, 18, 24])
>>> x2 = jnp.array([5, 10, 15])
>>> jnp.lcm(x1, x2)
Array([ 60,  90, 120], dtype=int32)

Broadcasting:

>>> x1 = jnp.array([12])
>>> x2 = jnp.array([6, 9, 12])
>>> jnp.lcm(x1, x2)
Array([12, 36, 12], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.ldexp(x1, x2, /)

Compute x1 * 2 ** x2

JAX implementation of numpy.ldexp().

Note that XLA does not provide an ldexp operation, so this is implemneted in JAX via a standard multiplication and exponentiation.

Args:

x1: real-valued input array. x2: integer input array. Must be broadcast-compatible with x1.

Returns:

x1 * 2 ** x2 computed element-wise.

See also:

  • jax.numpy.frexp(): decompose values into mantissa and exponent.

Examples:
>>> x1 = jnp.arange(5.0)
>>> x2 = 10
>>> jnp.ldexp(x1, x2)
Array([   0., 1024., 2048., 3072., 4096.], dtype=float32)

ldexp can be used to reconstruct the input to frexp:

>>> x = jnp.array([2., 3., 5., 11.])
>>> m, e = jnp.frexp(x)
>>> m
Array([0.5   , 0.75  , 0.625 , 0.6875], dtype=float32)
>>> e
Array([2, 2, 3, 4], dtype=int32)
>>> jnp.ldexp(m, e)
Array([ 2.,  3.,  5., 11.], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.left_shift(x, y, /)

Shift bits of x to left by the amount specified in y, element-wise.

JAX implementation of numpy.left_shift.

Args:

x: Input array, must be integer-typed. y: The amount of bits to shift each element in x to the left, only accepts

integer subtypes. x and y must either have same shape or be broadcast compatible.

Returns:

An array containing the left shifted elements of x by the amount specified in y, with the same shape as the broadcasted shape of x and y.

Note:

Left shifting x by y is equivalent to x * (2**y) within the bounds of the dtypes involved.

See also:

  • jax.numpy.right_shift(): and jax.numpy.bitwise_right_shift(): Shifts the bits of x1 to right by the amount specified in x2, element-wise.

  • jax.numpy.bitwise_left_shift(): Alias of jax.left_shift().

Examples:
>>> def print_binary(x):
...   return [bin(int(val)) for val in x]

>>> x1 = jnp.arange(5)
>>> x1
Array([0, 1, 2, 3, 4], dtype=int32)
>>> print_binary(x1)
['0b0', '0b1', '0b10', '0b11', '0b100']
>>> x2 = 1
>>> result = jnp.left_shift(x1, x2)
>>> result
Array([0, 2, 4, 6, 8], dtype=int32)
>>> print_binary(result)
['0b0', '0b10', '0b100', '0b110', '0b1000']

>>> x3 = 4
>>> print_binary([x3])
['0b100']
>>> x4 = jnp.array([1, 2, 3, 4])
>>> result1 = jnp.left_shift(x3, x4)
>>> result1
Array([ 8, 16, 32, 64], dtype=int32)
>>> print_binary(result1)
['0b1000', '0b10000', '0b100000', '0b1000000']
Parameters:
Return type:

Array

quchip.declarative.qnp.less(x, y, /)

Return element-wise truth value of x < y.

JAX implementation of numpy.less.

Args:

x: input array or scalar. y: input array or scalar. x and y must either have same shape or be

broadcast compatible.

Returns:

An array containing boolean values. True if the elements of x < y, and False otherwise.

See also:

  • jax.numpy.greater(): Returns element-wise truth value of x > y.

  • jax.numpy.greater_equal(): Returns element-wise truth value of x >= y.

  • jax.numpy.less_equal(): Returns element-wise truth value of x <= y.

Examples:

Scalar inputs:

>>> jnp.less(3, 7)
Array(True, dtype=bool, weak_type=True)

Inputs with same shape:

>>> x = jnp.array([5, 9, -3])
>>> y = jnp.array([1, 6, 4])
>>> jnp.less(x, y)
Array([False, False,  True], dtype=bool)

Inputs with broadcast compatibility:

>>> x1 = jnp.array([[2, -4, 6, -8],
...                 [-1, 5, -3, 7]])
>>> y1 = jnp.array([0, 3, -5, 9])
>>> jnp.less(x1, y1)
Array([[False,  True, False,  True],
       [ True, False, False,  True]], dtype=bool)
Parameters:
Return type:

Array

quchip.declarative.qnp.less_equal(x, y, /)

Return element-wise truth value of x <= y.

JAX implementation of numpy.less_equal.

Args:

x: input array or scalar. y: input array or scalar. x and y must have either same shape or be

broadcast compatible.

Returns:

An array containing the boolean values. True if the elements of x <= y, and False otherwise.

See also:

  • jax.numpy.greater_equal(): Returns element-wise truth value of x >= y.

  • jax.numpy.greater(): Returns element-wise truth value of x > y.

  • jax.numpy.less(): Returns element-wise truth value of x < y.

Examples:

Scalar inputs:

>>> jnp.less_equal(6, -2)
Array(False, dtype=bool, weak_type=True)

Inputs with same shape:

>>> x = jnp.array([-4, 1, 7])
>>> y = jnp.array([2, -3, 8])
>>> jnp.less_equal(x, y)
Array([ True, False,  True], dtype=bool)

Inputs with broadcast compatibility:

>>> x1 = jnp.array([2, -5, 9])
>>> y1 = jnp.array([[1, -6, 5],
...                 [-2, 4, -6]])
>>> jnp.less_equal(x1, y1)
Array([[False, False, False],
       [False,  True, False]], dtype=bool)
Parameters:
Return type:

Array

quchip.declarative.qnp.lexsort(keys, axis=-1)

Sort a sequence of keys in lexicographic order.

JAX implementation of numpy.lexsort().

Args:
keys: a sequence of arrays to sort; all arrays must have the same shape.

The last key in the sequence is used as the primary key.

axis: the axis along which to sort (default: -1).

Returns:

An array of integers of shape keys[0].shape giving the indices of the entries in lexicographically-sorted order.

See also:

  • jax.numpy.argsort(): sort a single entry by index.

  • jax.lax.sort(): direct XLA sorting API.

Examples:

lexsort() with a single key is equivalent to argsort():

>>> key1 = jnp.array([4, 2, 3, 2, 5])
>>> jnp.lexsort([key1])
Array([1, 3, 2, 0, 4], dtype=int32)
>>> jnp.argsort(key1)
Array([1, 3, 2, 0, 4], dtype=int32)

With multiple keys, lexsort() uses the last key as the primary key:

>>> key2 = jnp.array([2, 1, 1, 2, 2])
>>> jnp.lexsort([key1, key2])
Array([1, 2, 3, 0, 4], dtype=int32)

The meaning of the indices become more clear when printing the sorted keys:

>>> indices = jnp.lexsort([key1, key2])
>>> print(f"{key1[indices]}\n{key2[indices]}")
[2 3 2 4 5]
[1 1 2 2 2]

Notice that the elements of key2 appear in order, and within the sequences of duplicated values the corresponding elements of `key1 appear in order.

For multi-dimensional inputs, lexsort() defaults to sorting along the last axis:

>>> key1 = jnp.array([[2, 4, 2, 3],
...                   [3, 1, 2, 2]])
>>> key2 = jnp.array([[1, 2, 1, 3],
...                   [2, 1, 2, 1]])
>>> jnp.lexsort([key1, key2])
Array([[0, 2, 1, 3],
       [1, 3, 2, 0]], dtype=int32)

A different sort axis can be chosen using the axis keyword; here we sort along the leading axis:

>>> jnp.lexsort([key1, key2], axis=0)
Array([[0, 1, 0, 1],
       [1, 0, 1, 0]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None, axis=0, *, device=None)

Return evenly-spaced numbers within an interval.

JAX implementation of numpy.linspace().

Args:

start: scalar or array of starting values. stop: scalar or array of stop values. num: number of values to generate. Default: 50. endpoint: if True (default) then include the stop value in the result.

If False, then exclude the stop value.

retstep: If True, then return a (result, step) tuple, where step is the

interval between adjacent values in result.

axis: integer axis along which to generate the linspace. Defaults to zero. device: optional Device or Sharding

to which the created array will be committed.

Returns:

An array values, or a tuple (values, step) if retstep is True, where:

  • values is an array of evenly-spaced values from start to stop

  • step is the interval between adjacent values.

See also:

  • jax.numpy.arange(): Generate N evenly-spaced values given a starting point and a step

  • jax.numpy.logspace(): Generate logarithmically-spaced values.

  • jax.numpy.geomspace(): Generate geometrically-spaced values.

Examples:

List of 5 values between 0 and 10:

>>> jnp.linspace(0, 10, 5)
Array([ 0. ,  2.5,  5. ,  7.5, 10. ], dtype=float32)

List of 8 values between 0 and 10, excluding the endpoint:

>>> jnp.linspace(0, 10, 8, endpoint=False)
Array([0.  , 1.25, 2.5 , 3.75, 5.  , 6.25, 7.5 , 8.75], dtype=float32)

List of values and the step size between them

>>> vals, step = jnp.linspace(0, 10, 9, retstep=True)
>>> vals
Array([ 0.  ,  1.25,  2.5 ,  3.75,  5.  ,  6.25,  7.5 ,  8.75, 10.  ],      dtype=float32)
>>> step
Array(1.25, dtype=float32)

Multi-dimensional linspace:

>>> start = jnp.array([0, 5])
>>> stop = jnp.array([5, 10])
>>> jnp.linspace(start, stop, 5)
Array([[ 0.  ,  5.  ],
       [ 1.25,  6.25],
       [ 2.5 ,  7.5 ],
       [ 3.75,  8.75],
       [ 5.  , 10.  ]], dtype=float32)
Parameters:
Return type:

Array | tuple[Array, Array]

quchip.declarative.qnp.load(file, *args, **kwargs)

Load JAX arrays from npy files.

JAX wrapper of numpy.load().

This function is a simple wrapper of numpy.load(), but in the case of .npy files created with numpy.save() or jax.numpy.save(), the output will be returned as a jax.Array, and bfloat16 data types will be restored. For .npz files, results will be returned as normal NumPy arrays.

This function requires concrete array inputs, and is not compatible with transformations like jax.jit() or jax.vmap().

Args:

file: string, bytes, or path-like object containing the array data. args, kwargs: for additional arguments, see numpy.load()

Returns:

the array stored in the file.

See also:

  • jax.numpy.save(): save an array to a file.

Examples:
>>> import io
>>> f = io.BytesIO()  # use an in-memory file-like object.
>>> x = jnp.array([2, 4, 6, 8], dtype='bfloat16')
>>> jnp.save(f, x)
>>> f.seek(0)
0
>>> jnp.load(f)
Array([2, 4, 6, 8], dtype=bfloat16)
Parameters:
Return type:

Array

quchip.declarative.qnp.log(x, /)

Calculate element-wise natural logarithm of the input.

JAX implementation of numpy.log.

Args:

x: input array or scalar.

Returns:

An array containing the logarithm of each element in x, promotes to inexact dtype.

See also:

  • jax.numpy.exp(): Calculates element-wise exponential of the input.

  • jax.numpy.log2(): Calculates base-2 logarithm of each element of input.

  • jax.numpy.log1p(): Calculates element-wise logarithm of one plus input.

Examples:

jnp.log and jnp.exp are inverse functions of each other. Applying jnp.log on the result of jnp.exp(x) yields the original input x.

>>> x = jnp.array([2, 3, 4, 5])
>>> jnp.log(jnp.exp(x))
Array([2., 3., 4., 5.], dtype=float32)

Using jnp.log we can demonstrate well-known properties of logarithms, such as \(log(a*b) = log(a)+log(b)\).

>>> x1 = jnp.array([2, 1, 3, 1])
>>> x2 = jnp.array([1, 3, 2, 4])
>>> jnp.allclose(jnp.log(x1*x2), jnp.log(x1)+jnp.log(x2))
Array(True, dtype=bool)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.log10(x, /)

Calculates the base-10 logarithm of x element-wise

JAX implementation of numpy.log10.

Args:

x: Input array

Returns:

An array containing the base-10 logarithm of each element in x, promotes to inexact dtype.

Examples:
>>> x1 = jnp.array([0.01, 0.1, 1, 10, 100, 1000])
>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jnp.log10(x1))
[-2. -1.  0.  1.  2.  3.]
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.log1p(x, /)

Calculates element-wise logarithm of one plus input, log(x+1).

JAX implementation of numpy.log1p.

Args:

x: input array or scalar.

Returns:

An array containing the logarithm of one plus of each element in x, promotes to inexact dtype.

Note:

jnp.log1p is more accurate than when using the naive computation of log(x+1) for small values of x.

See also:

  • jax.numpy.expm1(): Calculates \(e^x-1\) of each element of the input.

  • jax.numpy.log2(): Calculates base-2 logarithm of each element of input.

  • jax.numpy.log(): Calculates element-wise logarithm of the input.

Examples:
>>> x = jnp.array([2, 5, 9, 4])
>>> jnp.allclose(jnp.log1p(x), jnp.log(x+1))
Array(True, dtype=bool)

For values very close to 0, jnp.log1p(x) is more accurate than jnp.log(x+1):

>>> x1 = jnp.array([1e-4, 1e-6, 2e-10])
>>> jnp.expm1(jnp.log1p(x1))
Array([1.00000005e-04, 9.99999997e-07, 2.00000003e-10], dtype=float32)
>>> jnp.expm1(jnp.log(x1+1))
Array([1.000166e-04, 9.536743e-07, 0.000000e+00], dtype=float32)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.log2(x, /)

Calculates the base-2 logarithm of x element-wise.

JAX implementation of numpy.log2.

Args:

x: Input array

Returns:

An array containing the base-2 logarithm of each element in x, promotes to inexact dtype.

Examples:
>>> x1 = jnp.array([0.25, 0.5, 1, 2, 4, 8])
>>> jnp.log2(x1)
Array([-2., -1.,  0.,  1.,  2.,  3.], dtype=float32)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.logical_not(x, /)

Compute NOT bool(x) element-wise.

JAX implementation of numpy.logical_not().

Args:

x: input array of any dtype.

Returns:

A boolean array that computes NOT bool(x) element-wise

See also:

  • jax.numpy.invert() or jax.numpy.bitwise_invert(): bitwise NOT operation

Examples:

Compute NOT x element-wise on a boolean array:

>>> x = jnp.array([True, False, True])
>>> jnp.logical_not(x)
Array([False,  True, False], dtype=bool)

For boolean input, this is equivalent to invert(), which implements the unary ~ operator:

>>> ~x
Array([False,  True, False], dtype=bool)

For non-boolean input, the input of logical_not() is implicitly cast to boolean:

>>> x = jnp.array([-1, 0, 1])
>>> jnp.logical_not(x)
Array([False,  True, False], dtype=bool)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None, axis=0)

Generate logarithmically-spaced values.

JAX implementation of numpy.logspace().

Args:
start: scalar or array. Used to specify the start value. The start value is

base ** start.

stop: scalar or array. Used to specify the stop value. The end value is

base ** stop.

num: int, optional, default=50. Number of values to generate. endpoint: bool, optional, default=True. If True, then include the stop value

in the result. If False, then exclude the stop value.

base: scalar or array, optional, default=10. Specifies the base of the logarithm. dtype: optional. Specifies the dtype of the output. axis: int, optional, default=0. Axis along which to generate the logspace.

Returns:

An array of logarithm.

See also:

  • jax.numpy.arange(): Generate N evenly-spaced values given a starting point and a step value.

  • jax.numpy.linspace(): Generate evenly-spaced values.

  • jax.numpy.geomspace(): Generate geometrically-spaced values.

Examples:

List 5 logarithmically spaced values between 1 (10 ** 0) and 100 (10 ** 2):

>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.logspace(0, 2, 5)
Array([  1.   ,   3.162,  10.   ,  31.623, 100.   ], dtype=float32)

List 5 logarithmically-spaced values between 1(10 ** 0) and 100 (10 ** 2), excluding endpoint:

>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.logspace(0, 2, 5, endpoint=False)
Array([ 1.   ,  2.512,  6.31 , 15.849, 39.811], dtype=float32)

List 7 logarithmically-spaced values between 1 (2 ** 0) and 4 (2 ** 2) with base 2:

>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.logspace(0, 2, 7, base=2)
Array([1.   , 1.26 , 1.587, 2.   , 2.52 , 3.175, 4.   ], dtype=float32)

Multi-dimensional logspace:

>>> start = jnp.array([0, 5])
>>> stop = jnp.array([5, 0])
>>> base = jnp.array([2, 3])
>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.logspace(start, stop, 5, base=base)
Array([[  1.   , 243.   ],
       [  2.378,  61.547],
       [  5.657,  15.588],
       [ 13.454,   3.948],
       [ 32.   ,   1.   ]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.mask_indices(n, mask_func, k=0, *, size=None)

Return indices of a mask of an (n, n) array.

Args:

n: static integer array dimension. mask_func: a function that takes a shape (n, n) array and

an optional offset k, and returns a shape (n, n) mask. Examples of functions with this signature are triu() and tril().

k: a scalar value passed to mask_func. size: optional argument specifying the static size of the output arrays.

This is passed to nonzero() when generating the indices from the mask.

Returns:

a tuple of indices where mask_func is nonzero.

See also:

  • jax.numpy.triu_indices(): compute mask_indices for triu().

  • jax.numpy.tril_indices(): compute mask_indices for tril().

Examples:

Calling mask_indices on built-in masking functions:

>>> jnp.mask_indices(3, jnp.triu)
(Array([0, 0, 0, 1, 1, 2], dtype=int32), Array([0, 1, 2, 1, 2, 2], dtype=int32))

>>> jnp.mask_indices(3, jnp.tril)
(Array([0, 1, 1, 2, 2, 2], dtype=int32), Array([0, 0, 1, 0, 1, 2], dtype=int32))

Calling mask_indices on a custom masking function:

>>> def mask_func(x, k=0):
...   i = jnp.arange(x.shape[0])[:, None]
...   j = jnp.arange(x.shape[1])
...   return (i + 1) % (j + 1 + k) == 0
>>> mask_func(jnp.ones((3, 3)))
Array([[ True, False, False],
       [ True,  True, False],
       [ True, False,  True]], dtype=bool)
>>> jnp.mask_indices(3, mask_func)
(Array([0, 1, 1, 2, 2], dtype=int32), Array([0, 0, 1, 0, 2], dtype=int32))
Parameters:
Return type:

tuple[Array, Array]

quchip.declarative.qnp.matmul(a, b, *, precision=None, preferred_element_type=None)

Perform a matrix multiplication.

JAX implementation of numpy.matmul().

Args:

a: first input array, of shape (N,) or (..., K, N). b: second input array. Must have shape (N,) or (..., N, M).

In the multi-dimensional case, leading dimensions must be broadcast-compatible with the leading dimensions of a.

precision: either None (default), which means the default precision for

the backend, a Precision enum value (Precision.DEFAULT, Precision.HIGH or Precision.HIGHEST) or a tuple of two such values indicating precision of a and b.

preferred_element_type: either None (default), which means the default

accumulation type for the input types, or a datatype, indicating to accumulate results to and return a result with that datatype.

Returns:

array containing the matrix product of the inputs. Shape is a.shape[:-1] if b.ndim == 1, otherwise the shape is (..., K, M), where leading dimensions of a and b are broadcast together.

See Also:

  • jax.numpy.linalg.vecdot(): batched vector product.

  • jax.numpy.linalg.tensordot(): batched tensor product.

  • jax.lax.dot_general(): general N-dimensional batched dot product.

Examples:

Vector dot products:

>>> a = jnp.array([1, 2, 3])
>>> b = jnp.array([4, 5, 6])
>>> jnp.matmul(a, b)
Array(32, dtype=int32)

Matrix dot product:

>>> a = jnp.array([[1, 2, 3],
...                [4, 5, 6]])
>>> b = jnp.array([[1, 2],
...                [3, 4],
...                [5, 6]])
>>> jnp.matmul(a, b)
Array([[22, 28],
       [49, 64]], dtype=int32)

For convenience, in all cases you can do the same computation using the @ operator:

>>> a @ b
Array([[22, 28],
       [49, 64]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.matrix_transpose(x, /)

Transpose the last two dimensions of an array.

JAX implementation of numpy.matrix_transpose(), implemented in terms of jax.lax.transpose().

Args:

x: input array, Must have x.ndim >= 2

Returns:

matrix-transposed copy of the array.

See Also:

  • jax.Array.mT: same operation accessed via an Array() property.

  • jax.numpy.transpose(): general multi-axis transpose

Note:

Unlike numpy.matrix_transpose(), jax.numpy.matrix_transpose() will return a copy rather than a view of the input array. However, under JIT, the compiler will optimize-away such copies when possible, so this doesn’t have performance impacts in practice.

Examples:

Here is a 2x2x2 matrix representing a batched 2x2 matrix:

>>> x = jnp.array([[[1, 2],
...                 [3, 4]],
...                [[5, 6],
...                 [7, 8]]])
>>> jnp.matrix_transpose(x)
Array([[[1, 3],
        [2, 4]],

       [[5, 7],
        [6, 8]]], dtype=int32)

For convenience, you can perform the same transpose via the mT property of jax.Array:

>>> x.mT
Array([[[1, 3],
        [2, 4]],

       [[5, 7],
        [6, 8]]], dtype=int32)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.matvec(x1, x2, /)

Batched matrix-vector product.

JAX implementation of numpy.matvec().

Args:

x1: array of shape (..., M, N) x2: array of shape (..., N). Leading dimensions must be broadcast-compatible

with leading dimensions of x1.

Returns:

An array of shape (..., M) containing the batched matrix-vector product.

See also:

  • jax.numpy.linalg.vecdot(): batched vector product.

  • jax.numpy.vecmat(): vector-matrix product.

  • jax.numpy.matmul(): general matrix multiplication.

Examples:

Simple matrix-vector product:

>>> x1 = jnp.array([[1, 2, 3],
...                 [4, 5, 6]])
>>> x2 = jnp.array([7, 8, 9])
>>> jnp.matvec(x1, x2)
Array([ 50, 122], dtype=int32)

Batched matrix-vector product:

>>> x2 = jnp.array([[7, 8, 9],
...                 [5, 6, 7]])
>>> jnp.matvec(x1, x2)
Array([[ 50, 122],
       [ 38,  92]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.max(a, axis=None, out=None, keepdims=False, initial=None, where=None)

Return the maximum of the array elements along a given axis.

JAX implementation of numpy.max().

Args:

a: Input array. axis: int or array, default=None. Axis along which the maximum to be computed.

If None, the maximum is computed along all the axes.

keepdims: bool, default=False. If true, reduced axes are left in the result

with size 1.

initial: int or array, default=None. Initial value for the maximum. where: int or array of boolean dtype, default=None. The elements to be used

in the maximum. Array should be broadcast compatible to the input. initial must be specified when where is used.

out: Unused by JAX.

Returns:

An array of maximum values along the given axis.

See also:

  • jax.numpy.min(): Compute the minimum of array elements along a given axis.

  • jax.numpy.sum(): Compute the sum of array elements along a given axis.

  • jax.numpy.prod(): Compute the product of array elements along a given axis.

Examples:

By default, jnp.max computes the maximum of elements along all the axes.

>>> x = jnp.array([[9, 3, 4, 5],
...                [5, 2, 7, 4],
...                [8, 1, 3, 6]])
>>> jnp.max(x)
Array(9, dtype=int32)

If axis=1, the maximum will be computed along axis 1.

>>> jnp.max(x, axis=1)
Array([9, 7, 8], dtype=int32)

If keepdims=True, ndim of the output will be same of that of the input.

>>> jnp.max(x, axis=1, keepdims=True)
Array([[9],
       [7],
       [8]], dtype=int32)

To include only specific elements in computing the maximum, you can use where. It can either have same dimension as input

>>> where=jnp.array([[0, 0, 1, 0],
...                  [0, 0, 1, 1],
...                  [1, 1, 1, 0]], dtype=bool)
>>> jnp.max(x, axis=1, keepdims=True, initial=0, where=where)
Array([[4],
       [7],
       [8]], dtype=int32)

or must be broadcast compatible with input.

>>> where = jnp.array([[False],
...                    [False],
...                    [False]])
>>> jnp.max(x, axis=0, keepdims=True, initial=0, where=where)
Array([[0, 0, 0, 0]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.mean(a, axis=None, dtype=None, out=None, keepdims=False, *, where=None)

Return the mean of array elements along a given axis.

JAX implementation of numpy.mean().

Args:

a: input array. axis: optional, int or sequence of ints, default=None. Axis along which the

mean to be computed. If None, mean is computed along all the axes.

dtype: The type of the output array. If None (default) then the output dtype

will be match the input dtype for floating point inputs, or be set to float32 or float64 for non-floating-point inputs.

keepdims: bool, default=False. If true, reduced axes are left in the result

with size 1.

where: optional, boolean array, default=None. The elements to be used in the

mean. Array should be broadcast compatible to the input.

out: Unused by JAX.

Returns:

An array of the mean along the given axis.

Notes:

For inputs of type float16 or bfloat16, the reductions will be performed at float32 precision.

See also:

  • jax.numpy.average(): Compute the weighted average of array elements

  • jax.numpy.sum(): Compute the sum of array elements.

Examples:

By default, the mean is computed along all the axes.

>>> x = jnp.array([[1, 3, 4, 2],
...                [5, 2, 6, 3],
...                [8, 1, 2, 9]])
>>> jnp.mean(x)
Array(3.8333335, dtype=float32)

If axis=1, the mean is computed along axis 1.

>>> jnp.mean(x, axis=1)
Array([2.5, 4. , 5. ], dtype=float32)

If keepdims=True, ndim of the output is equal to that of the input.

>>> jnp.mean(x, axis=1, keepdims=True)
Array([[2.5],
       [4. ],
       [5. ]], dtype=float32)

To use only specific elements of x to compute the mean, you can use where.

>>> where = jnp.array([[1, 0, 1, 0],
...                    [0, 1, 0, 1],
...                    [1, 1, 0, 1]], dtype=bool)
>>> jnp.mean(x, axis=1, keepdims=True, where=where)
Array([[2.5],
       [2.5],
       [6. ]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.median(a, axis=None, out=None, overwrite_input=False, keepdims=False)

Return the median of array elements along a given axis.

JAX implementation of numpy.median().

Args:

a: input array. axis: optional, int or sequence of ints, default=None. Axis along which the

median to be computed. If None, median is computed for the flattened array.

keepdims: bool, default=False. If true, reduced axes are left in the result

with size 1.

out: Unused by JAX. overwrite_input: Unused by JAX.

Returns:

An array of the median along the given axis.

See also:

  • jax.numpy.mean(): Compute the mean of array elements over a given axis.

  • jax.numpy.max(): Compute the maximum of array elements over given axis.

  • jax.numpy.min(): Compute the minimum of array elements over given axis.

Examples:

By default, the median is computed for the flattened array.

>>> x = jnp.array([[2, 4, 7, 1],
...                [3, 5, 9, 2],
...                [6, 1, 8, 3]])
>>> jnp.median(x)
Array(3.5, dtype=float32)

If axis=1, the median is computed along axis 1.

>>> jnp.median(x, axis=1)
Array([3. , 4. , 4.5], dtype=float32)

If keepdims=True, ndim of the output is equal to that of the input.

>>> jnp.median(x, axis=1, keepdims=True)
Array([[3. ],
       [4. ],
       [4.5]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.meshgrid(*xi, copy=True, sparse=False, indexing='xy')

Construct N-dimensional grid arrays from N 1-dimensional vectors.

JAX implementation of numpy.meshgrid().

Args:

xi: N arrays to convert to a grid. copy: whether to copy the input arrays. JAX supports only copy=True,

though under JIT compilation the compiler may opt to avoid copies.

sparse: if False (default), then each returned arrays will be of shape

[len(x1), len(x2), ..., len(xN)]. If False, then returned arrays will be of shape [1, 1, ..., len(xi), ..., 1, 1].

indexing: options are 'xy' for cartesian indexing (default) or 'ij'

for matrix indexing.

Returns:

A length-N list of grid arrays.

See also:

  • jax.numpy.indices(): generate a grid of indices.

  • jax.numpy.mgrid: create a meshgrid using indexing syntax.

  • jax.numpy.ogrid: create an open meshgrid using indexing syntax.

Examples:

For the following examples, we’ll use these 1D arrays as inputs:

>>> x = jnp.array([1, 2])
>>> y = jnp.array([10, 20, 30])

2D cartesian mesh grid:

>>> x_grid, y_grid = jnp.meshgrid(x, y)
>>> print(x_grid)
[[1 2]
 [1 2]
 [1 2]]
>>> print(y_grid)
[[10 10]
 [20 20]
 [30 30]]

2D sparse cartesian mesh grid:

>>> x_grid, y_grid = jnp.meshgrid(x, y, sparse=True)
>>> print(x_grid)
[[1 2]]
>>> print(y_grid)
[[10]
 [20]
 [30]]

2D matrix-index mesh grid:

>>> x_grid, y_grid = jnp.meshgrid(x, y, indexing='ij')
>>> print(x_grid)
[[1 1 1]
 [2 2 2]]
>>> print(y_grid)
[[10 20 30]
 [10 20 30]]
Parameters:
Return type:

list[Array]

quchip.declarative.qnp.min(a, axis=None, out=None, keepdims=False, initial=None, where=None)

Return the minimum of array elements along a given axis.

JAX implementation of numpy.min().

Args:

a: Input array. axis: int or array, default=None. Axis along which the minimum to be computed.

If None, the minimum is computed along all the axes.

keepdims: bool, default=False. If true, reduced axes are left in the result

with size 1.

initial: int or array, Default=None. Initial value for the minimum. where: int or array, default=None. The elements to be used in the minimum.

Array should be broadcast compatible to the input. initial must be specified when where is used.

out: Unused by JAX.

Returns:

An array of minimum values along the given axis.

See also:

  • jax.numpy.max(): Compute the maximum of array elements along a given axis.

  • jax.numpy.sum(): Compute the sum of array elements along a given axis.

  • jax.numpy.prod(): Compute the product of array elements along a given axis.

Examples:

By default, the minimum is computed along all the axes.

>>> x = jnp.array([[2, 5, 1, 6],
...                [3, -7, -2, 4],
...                [8, -4, 1, -3]])
>>> jnp.min(x)
Array(-7, dtype=int32)

If axis=1, the minimum is computed along axis 1.

>>> jnp.min(x, axis=1)
Array([ 1, -7, -4], dtype=int32)

If keepdims=True, ndim of the output will be same of that of the input.

>>> jnp.min(x, axis=1, keepdims=True)
Array([[ 1],
       [-7],
       [-4]], dtype=int32)

To include only specific elements in computing the minimum, you can use where. where can either have same dimension as input.

>>> where=jnp.array([[1, 0, 1, 0],
...                  [0, 0, 1, 1],
...                  [1, 1, 1, 0]], dtype=bool)
>>> jnp.min(x, axis=1, keepdims=True, initial=0, where=where)
Array([[ 0],
       [-2],
       [-4]], dtype=int32)

or must be broadcast compatible with input.

>>> where = jnp.array([[False],
...                    [False],
...                    [False]])
>>> jnp.min(x, axis=0, keepdims=True, initial=0, where=where)
Array([[0, 0, 0, 0]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.mod(x1, x2, /)

Alias of jax.numpy.remainder()

Parameters:
Return type:

Array

quchip.declarative.qnp.modf(x, /, out=None)

Return element-wise fractional and integral parts of the input array.

JAX implementation of numpy.modf.

Args:

x: input array or scalar. out: Not used by JAX.

Returns:

An array containing the fractional and integral parts of the elements of x, promoting dtypes inexact.

See also:

  • jax.numpy.divmod(): Calculates the integer quotient and remainder of x1 by x2 element-wise.

Examples:
>>> jnp.modf(4.8)
(Array(0.8000002, dtype=float32, weak_type=True), Array(4., dtype=float32, weak_type=True))
>>> x = jnp.array([-3.4, -5.7, 0.6, 1.5, 2.3])
>>> jnp.modf(x)
(Array([-0.4000001 , -0.6999998 ,  0.6       ,  0.5       ,  0.29999995],      dtype=float32), Array([-3., -5.,  0.,  1.,  2.], dtype=float32))
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

tuple[Array, Array]

quchip.declarative.qnp.moveaxis(a, source, destination)

Move an array axis to a new position

JAX implementation of numpy.moveaxis(), implemented in terms of jax.lax.transpose().

Args:

a: input array source: index or indices of the axes to move. destination: index or indices of the axes destinations

Returns:

Copy of a with axes moved from source to destination.

Notes:

Unlike numpy.moveaxis(), jax.numpy.moveaxis() will return a copy rather than a view of the input array. However, under JIT, the compiler will optimize away such copies when possible, so this doesn’t have performance impacts in practice.

See also:

  • jax.numpy.swapaxes(): swap two axes.

  • jax.numpy.rollaxis(): older API for moving an axis.

  • jax.numpy.transpose(): general axes permutation.

Examples:
>>> a = jnp.ones((2, 3, 4, 5))

Move axis 1 to the end of the array:

>>> jnp.moveaxis(a, 1, -1).shape
(2, 4, 5, 3)

Move the last axis to position 1:

>>> jnp.moveaxis(a, -1, 1).shape
(2, 5, 3, 4)

Move multiple axes:

>>> jnp.moveaxis(a, (0, 1), (-1, -2)).shape
(4, 5, 3, 2)

This can also be accomplished via transpose():

>>> a.transpose(2, 3, 1, 0).shape
(4, 5, 3, 2)
Parameters:
Return type:

Array

quchip.declarative.qnp.nan_to_num(x, copy=True, nan=0.0, posinf=None, neginf=None)

Replace NaN and infinite entries in an array.

JAX implementation of numpy.nan_to_num().

Args:
x: array of values to be replaced. If it does not have an inexact

dtype it will be returned unmodified.

copy: unused by JAX nan: value to substitute for NaN entries. Defaults to 0.0. posinf: value to substitute for positive infinite entries.

Defaults to the maximum representable value.

neginf: value to substitute for positive infinite entries.

Defaults to the minimum representable value.

Returns:

A copy of x with the requested substitutions.

See also:

  • jax.numpy.isnan(): return True where the array contains NaN

  • jax.numpy.isposinf(): return True where the array contains +inf

  • jax.numpy.isneginf(): return True where the array contains -inf

Examples:
>>> x = jnp.array([0, jnp.nan, 1, jnp.inf, 2, -jnp.inf])

Default substitution values:

>>> jnp.nan_to_num(x)
Array([ 0.0000000e+00,  0.0000000e+00,  1.0000000e+00,  3.4028235e+38,
        2.0000000e+00, -3.4028235e+38], dtype=float32)

Overriding substitutions for -inf and +inf:

>>> jnp.nan_to_num(x, posinf=999, neginf=-999)
Array([   0.,    0.,    1.,  999.,    2., -999.], dtype=float32)

If you only wish to substitute for NaN values while leaving inf values untouched, using where() with jax.numpy.isnan() is a better option:

>>> jnp.where(jnp.isnan(x), 0, x)
Array([  0.,   0.,   1.,  inf,   2., -inf], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.nanargmax(a, axis=None, out=None, keepdims=None)

Return the index of the maximum value of an array, ignoring NaNs.

JAX implementation of numpy.nanargmax().

Args:

a: input array axis: optional integer specifying the axis along which to find the maximum

value. If axis is not specified, a will be flattened.

out: unused by JAX keepdims: if True, then return an array with the same number of dimensions

as a.

Returns:

an array containing the index of the maximum value along the specified axis.

Note:

In the case of an axis with all-NaN values, the returned index will be -1. This differs from the behavior of numpy.nanargmax(), which raises an error.

See also:

  • jax.numpy.argmax(): return the index of the maximum value.

  • jax.numpy.nanargmin(): compute argmin while ignoring NaN values.

Examples:
>>> x = jnp.array([1, 3, 5, 4, jnp.nan])

Using a standard argmax() leads to potentially unexpected results:

>>> jnp.argmax(x)
Array(4, dtype=int32)

Using nanargmax returns the index of the maximum non-NaN value.

>>> jnp.nanargmax(x)
Array(2, dtype=int32)

>>> x = jnp.array([[1, 3, jnp.nan],
...                [5, 4, jnp.nan]])
>>> jnp.nanargmax(x, axis=1)
Array([1, 0], dtype=int32)

>>> jnp.nanargmax(x, axis=1, keepdims=True)
Array([[1],
       [0]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.nanargmin(a, axis=None, out=None, keepdims=None)

Return the index of the minimum value of an array, ignoring NaNs.

JAX implementation of numpy.nanargmin().

Args:

a: input array axis: optional integer specifying the axis along which to find the maximum

value. If axis is not specified, a will be flattened.

out: unused by JAX keepdims: if True, then return an array with the same number of dimensions

as a.

Returns:

an array containing the index of the minimum value along the specified axis.

Note:

In the case of an axis with all-NaN values, the returned index will be -1. This differs from the behavior of numpy.nanargmin(), which raises an error.

See also:

  • jax.numpy.argmin(): return the index of the minimum value.

  • jax.numpy.nanargmax(): compute argmax while ignoring NaN values.

Examples:
>>> x = jnp.array([jnp.nan, 3, 5, 4, 2])
>>> jnp.nanargmin(x)
Array(4, dtype=int32)

>>> x = jnp.array([[1, 3, jnp.nan],
...                [5, 4, jnp.nan]])
>>> jnp.nanargmin(x, axis=1)
Array([0, 1], dtype=int32)

>>> jnp.nanargmin(x, axis=1, keepdims=True)
Array([[0],
       [1]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.nancumprod(a, axis=None, dtype=None, out=None)

Cumulative product of elements along an axis, ignoring NaN values.

JAX implementation of numpy.nancumprod().

Args:

a: N-dimensional array to be accumulated. axis: integer axis along which to accumulate. If None (default), then

array will be flattened and accumulated along the flattened axis.

dtype: optionally specify the dtype of the output. If not specified,

then the output dtype will match the input dtype.

out: unused by JAX

Returns:

An array containing the accumulated product along the given axis.

See also:

  • jax.numpy.cumprod(): cumulative product without ignoring NaN values.

  • jax.numpy.multiply.accumulate(): cumulative product via ufunc methods.

  • jax.numpy.prod(): product along axis

Examples:
>>> x = jnp.array([[1., 2., jnp.nan],
...                [4., jnp.nan, 6.]])

The standard cumulative product will propagate NaN values:

>>> jnp.cumprod(x)
Array([ 1.,  2., nan, nan, nan, nan], dtype=float32)

nancumprod() will ignore NaN values, effectively replacing them with ones:

>>> jnp.nancumprod(x)
Array([ 1.,  2.,  2.,  8.,  8., 48.], dtype=float32)

Cumulative product along axis 1:

>>> jnp.nancumprod(x, axis=1)
Array([[ 1.,  2.,  2.],
       [ 4.,  4., 24.]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.nancumsum(a, axis=None, dtype=None, out=None)

Cumulative sum of elements along an axis, ignoring NaN values.

JAX implementation of numpy.nancumsum().

Args:

a: N-dimensional array to be accumulated. axis: integer axis along which to accumulate. If None (default), then

array will be flattened and accumulated along the flattened axis.

dtype: optionally specify the dtype of the output. If not specified,

then the output dtype will match the input dtype.

out: unused by JAX

Returns:

An array containing the accumulated sum along the given axis.

See also:

  • jax.numpy.cumsum(): cumulative sum without ignoring NaN values.

  • jax.numpy.cumulative_sum(): cumulative sum via the array API standard.

  • jax.numpy.add.accumulate(): cumulative sum via ufunc methods.

  • jax.numpy.sum(): sum along axis

Examples:
>>> x = jnp.array([[1., 2., jnp.nan],
...                [4., jnp.nan, 6.]])

The standard cumulative sum will propagate NaN values:

>>> jnp.cumsum(x)
Array([ 1.,  3., nan, nan, nan, nan], dtype=float32)

nancumsum() will ignore NaN values, effectively replacing them with zeros:

>>> jnp.nancumsum(x)
Array([ 1.,  3.,  3.,  7.,  7., 13.], dtype=float32)

Cumulative sum along axis 1:

>>> jnp.nancumsum(x, axis=1)
Array([[ 1.,  3.,  3.],
       [ 4.,  4., 10.]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.nanmax(a, axis=None, out=None, keepdims=False, initial=None, where=None)

Return the maximum of the array elements along a given axis, ignoring NaNs.

JAX implementation of numpy.nanmax().

Args:

a: Input array. axis: int or sequence of ints, default=None. Axis along which the maximum is

computed. If None, the maximum is computed along the flattened array.

keepdims: bool, default=False. If True, reduced axes are left in the result

with size 1.

initial: int or array, default=None. Initial value for the maximum. where: array of boolean dtype, default=None. The elements to be used in the

maximum. Array should be broadcast compatible to the input. initial must be specified when where is used.

out: Unused by JAX.

Returns:

An array of maximum values along the given axis, ignoring NaNs. If all values are NaNs along the given axis, returns nan.

See also:

  • jax.numpy.nanmin(): Compute the minimum of array elements along a given axis, ignoring NaNs.

  • jax.numpy.nansum(): Compute the sum of array elements along a given axis, ignoring NaNs.

  • jax.numpy.nanprod(): Compute the product of array elements along a given axis, ignoring NaNs.

  • jax.numpy.nanmean(): Compute the mean of array elements along a given axis, ignoring NaNs.

Examples:

By default, jnp.nanmax computes the maximum of elements along the flattened array.

>>> nan = jnp.nan
>>> x = jnp.array([[8, nan, 4, 6],
...                [nan, -2, nan, -4],
...                [-2, 1, 7, nan]])
>>> jnp.nanmax(x)
Array(8., dtype=float32)

If axis=1, the maximum will be computed along axis 1.

>>> jnp.nanmax(x, axis=1)
Array([ 8., -2.,  7.], dtype=float32)

If keepdims=True, ndim of the output will be same of that of the input.

>>> jnp.nanmax(x, axis=1, keepdims=True)
Array([[ 8.],
       [-2.],
       [ 7.]], dtype=float32)

To include only specific elements in computing the maximum, you can use where. It can either have same dimension as input

>>> where=jnp.array([[0, 0, 1, 0],
...                  [0, 0, 1, 1],
...                  [1, 1, 1, 0]], dtype=bool)
>>> jnp.nanmax(x, axis=1, keepdims=True, initial=0, where=where)
Array([[4.],
       [0.],
       [7.]], dtype=float32)

or must be broadcast compatible with input.

>>> where = jnp.array([[True],
...                    [False],
...                    [False]])
>>> jnp.nanmax(x, axis=0, keepdims=True, initial=0, where=where)
Array([[8., 0., 4., 6.]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.nanmean(a, axis=None, dtype=None, out=None, keepdims=False, where=None)

Return the mean of the array elements along a given axis, ignoring NaNs.

JAX implementation of numpy.nanmean().

Args:

a: Input array. axis: int or sequence of ints, default=None. Axis along which the mean is

computed. If None, the mean is computed along the flattened array.

dtype: The type of the output array. Default=None. keepdims: bool, default=False. If True, reduced axes are left in the result

with size 1.

where: array of boolean dtype, default=None. The elements to be used in

computing mean. Array should be broadcast compatible to the input.

out: Unused by JAX.

Returns:

An array containing the mean of array elements along the given axis, ignoring NaNs. If all elements along the given axis are NaNs, returns nan.

See also:

  • jax.numpy.nanmin(): Compute the minimum of array elements along a given axis, ignoring NaNs.

  • jax.numpy.nanmax(): Compute the maximum of array elements along a given axis, ignoring NaNs.

  • jax.numpy.nansum(): Compute the sum of array elements along a given axis, ignoring NaNs.

  • jax.numpy.nanprod(): Compute the product of array elements along a given axis, ignoring NaNs.

Examples:

By default, jnp.nanmean computes the mean of elements along the flattened array.

>>> nan = jnp.nan
>>> x = jnp.array([[2, nan, 4, 3],
...                [nan, -2, nan, 9],
...                [4, -7, 6, nan]])
>>> jnp.nanmean(x)
Array(2.375, dtype=float32)

If axis=1, mean will be computed along axis 1.

>>> jnp.nanmean(x, axis=1)
Array([3. , 3.5, 1. ], dtype=float32)

If keepdims=True, ndim of the output will be same of that of the input.

>>> jnp.nanmean(x, axis=1, keepdims=True)
Array([[3. ],
       [3.5],
       [1. ]], dtype=float32)

where can be used to include only specific elements in computing the mean.

>>> where = jnp.array([[1, 0, 1, 0],
...                    [0, 0, 1, 1],
...                    [1, 1, 0, 1]], dtype=bool)
>>> jnp.nanmean(x, axis=1, keepdims=True, where=where)
Array([[ 3. ],
       [ 9. ],
       [-1.5]], dtype=float32)

If where is False at all elements, jnp.nanmean returns nan along the given axis.

>>> where = jnp.array([[False],
...                    [False],
...                    [False]])
>>> jnp.nanmean(x, axis=0, keepdims=True, where=where)
Array([[nan, nan, nan, nan]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.nanmedian(a, axis=None, out=None, overwrite_input=False, keepdims=False)

Return the median of array elements along a given axis, ignoring NaNs.

JAX implementation of numpy.nanmedian().

Args:

a: input array. axis: optional, int or sequence of ints, default=None. Axis along which the

median to be computed. If None, median is computed for the flattened array.

keepdims: bool, default=False. If true, reduced axes are left in the result

with size 1.

out: Unused by JAX. overwrite_input: Unused by JAX.

Returns:

An array containing the median along the given axis, ignoring NaNs. If all elements along the given axis are NaNs, returns nan.

See also:

  • jax.numpy.nanmean(): Compute the mean of array elements over a given axis, ignoring NaNs.

  • jax.numpy.nanmax(): Compute the maximum of array elements over given axis, ignoring NaNs.

  • jax.numpy.nanmin(): Compute the minimum of array elements over given axis, ignoring NaNs.

Examples:

By default, the median is computed for the flattened array.

>>> nan = jnp.nan
>>> x = jnp.array([[2, nan, 7, nan],
...                [nan, 5, 9, 2],
...                [6, 1, nan, 3]])
>>> jnp.nanmedian(x)
Array(4., dtype=float32)

If axis=1, the median is computed along axis 1.

>>> jnp.nanmedian(x, axis=1)
Array([4.5, 5. , 3. ], dtype=float32)

If keepdims=True, ndim of the output is equal to that of the input.

>>> jnp.nanmedian(x, axis=1, keepdims=True)
Array([[4.5],
       [5. ],
       [3. ]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.nanmin(a, axis=None, out=None, keepdims=False, initial=None, where=None)

Return the minimum of the array elements along a given axis, ignoring NaNs.

JAX implementation of numpy.nanmin().

Args:

a: Input array. axis: int or sequence of ints, default=None. Axis along which the minimum is

computed. If None, the minimum is computed along the flattened array.

keepdims: bool, default=False. If True, reduced axes are left in the result

with size 1.

initial: int or array, default=None. Initial value for the minimum. where: array of boolean dtype, default=None. The elements to be used in the

minimum. Array should be broadcast compatible to the input. initial must be specified when where is used.

out: Unused by JAX.

Returns:

An array of minimum values along the given axis, ignoring NaNs. If all values are NaNs along the given axis, returns nan.

See also:

  • jax.numpy.nanmax(): Compute the maximum of array elements along a given axis, ignoring NaNs.

  • jax.numpy.nansum(): Compute the sum of array elements along a given axis, ignoring NaNs.

  • jax.numpy.nanprod(): Compute the product of array elements along a given axis, ignoring NaNs.

  • jax.numpy.nanmean(): Compute the mean of array elements along a given axis, ignoring NaNs.

Examples:

By default, jnp.nanmin computes the minimum of elements along the flattened array.

>>> nan = jnp.nan
>>> x = jnp.array([[1, nan, 4, 5],
...                [nan, -2, nan, -4],
...                [2, 1, 3, nan]])
>>> jnp.nanmin(x)
Array(-4., dtype=float32)

If axis=1, the maximum will be computed along axis 1.

>>> jnp.nanmin(x, axis=1)
Array([ 1., -4.,  1.], dtype=float32)

If keepdims=True, ndim of the output will be same of that of the input.

>>> jnp.nanmin(x, axis=1, keepdims=True)
Array([[ 1.],
       [-4.],
       [ 1.]], dtype=float32)

To include only specific elements in computing the maximum, you can use where. It can either have same dimension as input

>>> where=jnp.array([[0, 0, 1, 0],
...                  [0, 0, 1, 1],
...                  [1, 1, 1, 0]], dtype=bool)
>>> jnp.nanmin(x, axis=1, keepdims=True, initial=0, where=where)
Array([[ 0.],
       [-4.],
       [ 0.]], dtype=float32)

or must be broadcast compatible with input.

>>> where = jnp.array([[False],
...                    [True],
...                    [False]])
>>> jnp.nanmin(x, axis=0, keepdims=True, initial=0, where=where)
Array([[ 0., -2.,  0., -4.]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.nanpercentile(a, q, axis=None, out=None, overwrite_input=False, method='linear', keepdims=False, *, interpolation=Deprecated)

Compute the percentile of the data along the specified axis, ignoring NaN values.

JAX implementation of numpy.nanpercentile().

Args:

a: N-dimensional array input. q: scalar or 1-dimensional array specifying the desired quantiles. q

should contain integer or floating point values between 0 and 100.

axis: optional axis or tuple of axes along which to compute the quantile out: not implemented by JAX; will error if not None overwrite_input: not implemented by JAX; will error if not False method: specify the interpolation method to use. Options are one of

["linear", "lower", "higher", "midpoint", "nearest"]. default is linear.

keepdims: if True, then the returned array will have the same number of

dimensions as the input. Default is False.

interpolation: deprecated alias of the method argument. Will result

in a DeprecationWarning if used.

Returns:

An array containing the specified percentiles along the specified axes.

See also:

  • jax.numpy.nanquantile(): compute the nan-aware quantile (0.0-1.0)

  • jax.numpy.percentile(): compute the percentile without special handling of NaNs.

Examples:

Computing the median and quartiles of a 1D array:

>>> x = jnp.array([0, 1, 2, jnp.nan, 3, 4, 5, 6])
>>> q = jnp.array([25, 50, 75])

Because of the NaN value, jax.numpy.percentile() returns all NaNs, while nanpercentile() ignores them:

>>> jnp.percentile(x, q)
Array([nan, nan, nan], dtype=float32)
>>> jnp.nanpercentile(x, q)
Array([1.5, 3. , 4.5], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.nanprod(a, axis=None, dtype=None, out=None, keepdims=False, initial=None, where=None)

Return the product of the array elements along a given axis, ignoring NaNs.

JAX implementation of numpy.nanprod().

Args:

a: Input array. axis: int or sequence of ints, default=None. Axis along which the product is

computed. If None, the product is computed along the flattened array.

dtype: The type of the output array. Default=None. keepdims: bool, default=False. If True, reduced axes are left in the result

with size 1.

initial: int or array, default=None. Initial value for the product. where: array of boolean dtype, default=None. The elements to be used in the

product. Array should be broadcast compatible to the input.

out: Unused by JAX.

Returns:

An array containing the product of array elements along the given axis, ignoring NaNs. If all elements along the given axis are NaNs, returns 1.

See also:

  • jax.numpy.nanmin(): Compute the minimum of array elements along a given axis, ignoring NaNs.

  • jax.numpy.nanmax(): Compute the maximum of array elements along a given axis, ignoring NaNs.

  • jax.numpy.nansum(): Compute the sum of array elements along a given axis, ignoring NaNs.

  • jax.numpy.nanmean(): Compute the mean of array elements along a given axis, ignoring NaNs.

Examples:

By default, jnp.nanprod computes the product of elements along the flattened array.

>>> nan = jnp.nan
>>> x = jnp.array([[nan, 3, 4, nan],
...                [5, nan, 1, 3],
...                [2, 1, nan, 1]])
>>> jnp.nanprod(x)
Array(360., dtype=float32)

If axis=1, the product will be computed along axis 1.

>>> jnp.nanprod(x, axis=1)
Array([12., 15.,  2.], dtype=float32)

If keepdims=True, ndim of the output will be same of that of the input.

>>> jnp.nanprod(x, axis=1, keepdims=True)
Array([[12.],
       [15.],
       [ 2.]], dtype=float32)

To include only specific elements in computing the maximum, you can use where.

>>> where=jnp.array([[1, 0, 1, 0],
...                  [0, 0, 1, 1],
...                  [1, 1, 1, 0]], dtype=bool)
>>> jnp.nanprod(x, axis=1, keepdims=True, where=where)
Array([[4.],
       [3.],
       [2.]], dtype=float32)

If where is False at all elements, jnp.nanprod returns 1 along the given axis.

>>> where = jnp.array([[False],
...                    [False],
...                    [False]])
>>> jnp.nanprod(x, axis=0, keepdims=True, where=where)
Array([[1., 1., 1., 1.]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.nanquantile(a, q, axis=None, out=None, overwrite_input=False, method='linear', keepdims=False, *, interpolation=Deprecated)

Compute the quantile of the data along the specified axis, ignoring NaNs.

JAX implementation of numpy.nanquantile().

Args:

a: N-dimensional array input. q: scalar or 1-dimensional array specifying the desired quantiles. q

should contain floating-point values between 0.0 and 1.0.

axis: optional axis or tuple of axes along which to compute the quantile out: not implemented by JAX; will error if not None overwrite_input: not implemented by JAX; will error if not False method: specify the interpolation method to use. Options are one of

["linear", "lower", "higher", "midpoint", "nearest"]. default is linear.

keepdims: if True, then the returned array will have the same number of

dimensions as the input. Default is False.

interpolation: deprecated alias of the method argument. Will result

in a DeprecationWarning if used.

Returns:

An array containing the specified quantiles along the specified axes.

See also:

  • jax.numpy.quantile(): compute the quantile without ignoring nans

  • jax.numpy.nanpercentile(): compute the percentile (0-100)

Examples:

Computing the median and quartiles of a 1D array:

>>> x = jnp.array([0, 1, 2, jnp.nan, 3, 4, 5, 6])
>>> q = jnp.array([0.25, 0.5, 0.75])

Because of the NaN value, jax.numpy.quantile() returns all NaNs, while nanquantile() ignores them:

>>> jnp.quantile(x, q)
Array([nan, nan, nan], dtype=float32)
>>> jnp.nanquantile(x, q)
Array([1.5, 3. , 4.5], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.nanstd(a, axis=None, dtype=None, out=None, ddof=0, keepdims=False, where=None)

Compute the standard deviation along a given axis, ignoring NaNs.

JAX implementation of numpy.nanstd().

Args:

a: input array. axis: optional, int or sequence of ints, default=None. Axis along which the

standard deviation is computed. If None, standard deviaiton is computed along flattened array.

dtype: The type of the output array. Default=None. ddof: int, default=0. Degrees of freedom. The divisor in the standard deviation

computation is N-ddof, N is number of elements along given axis.

keepdims: bool, default=False. If true, reduced axes are left in the result

with size 1.

where: optional, boolean array, default=None. The elements to be used in the

standard deviation. Array should be broadcast compatible to the input.

out: Unused by JAX.

Returns:

An array containing the standard deviation of array elements along the given axis. If all elements along the given axis are NaNs, returns nan.

See also:

  • jax.numpy.nanmean(): Compute the mean of array elements over a given axis, ignoring NaNs.

  • jax.numpy.nanvar(): Compute the variance along the given axis, ignoring NaNs values.

  • jax.numpy.std(): Computed the standard deviation along the given axis.

Examples:

By default, jnp.nanstd computes the standard deviation along flattened array.

>>> nan = jnp.nan
>>> x = jnp.array([[3, nan, 4, 5],
...                [nan, 2, nan, 7],
...                [2, 1, 6, nan]])
>>> jnp.nanstd(x)
Array(1.9843135, dtype=float32)

If axis=0, computes standard deviation along axis 0.

>>> jnp.nanstd(x, axis=0)
Array([0.5, 0.5, 1. , 1. ], dtype=float32)

To preserve the dimensions of input, you can set keepdims=True.

>>> jnp.nanstd(x, axis=0, keepdims=True)
Array([[0.5, 0.5, 1. , 1. ]], dtype=float32)

If ddof=1:

>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jnp.nanstd(x, axis=0, keepdims=True, ddof=1))
[[0.71 0.71 1.41 1.41]]

To include specific elements of the array to compute standard deviation, you can use where.

>>> where=jnp.array([[1, 0, 1, 0],
...                  [0, 1, 0, 1],
...                  [1, 1, 0, 1]], dtype=bool)
>>> jnp.nanstd(x, axis=0, keepdims=True, where=where)
Array([[0.5, 0.5, 0. , 0. ]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.nansum(a, axis=None, dtype=None, out=None, keepdims=False, initial=None, where=None)

Return the sum of the array elements along a given axis, ignoring NaNs.

JAX implementation of numpy.nansum().

Args:

a: Input array. axis: int or sequence of ints, default=None. Axis along which the sum is

computed. If None, the sum is computed along the flattened array.

dtype: The type of the output array. Default=None. keepdims: bool, default=False. If True, reduced axes are left in the result

with size 1.

initial: int or array, default=None. Initial value for the sum. where: array of boolean dtype, default=None. The elements to be used in the

sum. Array should be broadcast compatible to the input.

out: Unused by JAX.

Returns:

An array containing the sum of array elements along the given axis, ignoring NaNs. If all elements along the given axis are NaNs, returns 0.

See also:

  • jax.numpy.nanmin(): Compute the minimum of array elements along a given axis, ignoring NaNs.

  • jax.numpy.nanmax(): Compute the maximum of array elements along a given axis, ignoring NaNs.

  • jax.numpy.nanprod(): Compute the product of array elements along a given axis, ignoring NaNs.

  • jax.numpy.nanmean(): Compute the mean of array elements along a given axis, ignoring NaNs.

Examples:

By default, jnp.nansum computes the sum of elements along the flattened array.

>>> nan = jnp.nan
>>> x = jnp.array([[3, nan, 4, 5],
...                [nan, -2, nan, 7],
...                [2, 1, 6, nan]])
>>> jnp.nansum(x)
Array(26., dtype=float32)

If axis=1, the sum will be computed along axis 1.

>>> jnp.nansum(x, axis=1)
Array([12.,  5.,  9.], dtype=float32)

If keepdims=True, ndim of the output will be same of that of the input.

>>> jnp.nansum(x, axis=1, keepdims=True)
Array([[12.],
       [ 5.],
       [ 9.]], dtype=float32)

To include only specific elements in computing the sum, you can use where.

>>> where=jnp.array([[1, 0, 1, 0],
...                  [0, 0, 1, 1],
...                  [1, 1, 1, 0]], dtype=bool)
>>> jnp.nansum(x, axis=1, keepdims=True, where=where)
Array([[7.],
       [7.],
       [9.]], dtype=float32)

If where is False at all elements, jnp.nansum returns 0 along the given axis.

>>> where = jnp.array([[False],
...                    [False],
...                    [False]])
>>> jnp.nansum(x, axis=0, keepdims=True, where=where)
Array([[0., 0., 0., 0.]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.nanvar(a, axis=None, dtype=None, out=None, ddof=0, keepdims=False, where=None)

Compute the variance of array elements along a given axis, ignoring NaNs.

JAX implementation of numpy.nanvar().

Args:

a: input array. axis: optional, int or sequence of ints, default=None. Axis along which the

variance is computed. If None, variance is computed along flattened array.

dtype: The type of the output array. Default=None. ddof: int, default=0. Degrees of freedom. The divisor in the variance computation

is N-ddof, N is number of elements along given axis.

keepdims: bool, default=False. If true, reduced axes are left in the result

with size 1.

where: optional, boolean array, default=None. The elements to be used in the

variance. Array should be broadcast compatible to the input.

out: Unused by JAX.

Returns:

An array containing the variance of array elements along specified axis. If all elements along the given axis are NaNs, returns nan.

See also:

  • jax.numpy.nanmean(): Compute the mean of array elements over a given axis, ignoring NaNs.

  • jax.numpy.nanstd(): Computed the standard deviation of a given axis, ignoring NaNs.

  • jax.numpy.var(): Compute the variance of array elements along a given axis.

Examples:

By default, jnp.nanvar computes the variance along all axes.

>>> nan = jnp.nan
>>> x = jnp.array([[1, nan, 4, 3],
...                [nan, 2, nan, 9],
...                [4, 8, 6, nan]])
>>> jnp.nanvar(x)
Array(6.984375, dtype=float32)

If axis=1, variance is computed along axis 1.

>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jnp.nanvar(x, axis=1))
[ 1.56 12.25  2.67]

To preserve the dimensions of input, you can set keepdims=True.

>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jnp.nanvar(x, axis=1, keepdims=True))
[[ 1.56]
 [12.25]
 [ 2.67]]

If ddof=1:

>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jnp.nanvar(x, axis=1, keepdims=True, ddof=1))
[[ 2.33]
 [24.5 ]
 [ 4.  ]]

To include specific elements of the array to compute variance, you can use where.

>>> where = jnp.array([[1, 0, 1, 0],
...                    [0, 1, 1, 0],
...                    [1, 1, 0, 1]], dtype=bool)
>>> jnp.nanvar(x, axis=1, keepdims=True, where=where)
Array([[2.25],
       [0.  ],
       [4.  ]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.ndarray

alias of Array

quchip.declarative.qnp.ndim(a)

Return the number of dimensions of an array.

JAX implementation of numpy.ndim(). Unlike np.ndim, this function raises a TypeError if the input is a collection such as a list or tuple.

Args:

a: array-like object, or any object with an ndim attribute.

Returns:

An integer specifying the number of dimensions of a.

Examples:

Number of dimensions for arrays:

>>> x = jnp.arange(10)
>>> jnp.ndim(x)
1
>>> y = jnp.ones((2, 3))
>>> jnp.ndim(y)
2

This also works for scalars:

>>> jnp.ndim(3.14)
0

For arrays, this can also be accessed via the jax.Array.ndim property:

>>> x.ndim
1
Parameters:

a (Array | ndarray | bool | number | bool | int | float | complex | SupportsNdim)

Return type:

int

quchip.declarative.qnp.nextafter(x, y, /)

Return element-wise next floating point value after x towards y.

JAX implementation of numpy.nextafter.

Args:

x: scalar or array. Specifies the value after which the next number is found. y: scalar or array. Specifies the direction towards which the next number is

found. x and y should either have same shape or be broadcast compatible.

Returns:

An array containing the next representable number of x in the direction of y.

Examples:
>>> jnp.nextafter(2, 1)
Array(1.9999999, dtype=float32, weak_type=True)
>>> x = jnp.array([3, -2, 1])
>>> y = jnp.array([2, -1, 2])
>>> jnp.nextafter(x, y)
Array([ 2.9999998, -1.9999999,  1.0000001], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.nonzero(a, *, size=None, fill_value=None)

Return indices of nonzero elements of an array.

JAX implementation of numpy.nonzero().

Because the size of the output of nonzero is data-dependent, the function is not compatible with JIT and other transformations. The JAX version adds the optional size argument which must be specified statically for jnp.nonzero to be used within JAX’s transformations.

Args:

a: N-dimensional array. size: optional static integer specifying the number of nonzero entries to

return. If there are more nonzero elements than the specified size, then indices will be truncated at the end. If there are fewer nonzero elements than the specified size, then indices will be padded with fill_value, which defaults to zero.

fill_value: optional padding value when size is specified. Defaults to 0.

Returns:

Tuple of JAX Arrays of length a.ndim, containing the indices of each nonzero value.

See also:

  • jax.numpy.flatnonzero()

  • jax.numpy.where()

Examples:

One-dimensional array returns a length-1 tuple of indices:

>>> x = jnp.array([0, 5, 0, 6, 0, 7])
>>> jnp.nonzero(x)
(Array([1, 3, 5], dtype=int32),)

Two-dimensional array returns a length-2 tuple of indices:

>>> x = jnp.array([[0, 5, 0],
...                [6, 0, 7]])
>>> jnp.nonzero(x)
(Array([0, 1, 1], dtype=int32), Array([1, 0, 2], dtype=int32))

In either case, the resulting tuple of indices can be used directly to extract the nonzero values:

>>> indices = jnp.nonzero(x)
>>> x[indices]
Array([5, 6, 7], dtype=int32)

The output of nonzero has a dynamic shape, because the number of returned indices depends on the contents of the input array. As such, it is incompatible with JIT and other JAX transformations:

>>> x = jnp.array([0, 5, 0, 6, 0, 7])
>>> jax.jit(jnp.nonzero)(x)
Traceback (most recent call last):
  ...
ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: traced array with shape int32[].
The size argument of jnp.nonzero must be statically specified to use jnp.nonzero within JAX transformations.

This can be addressed by passing a static size parameter to specify the desired output shape:

>>> nonzero_jit = jax.jit(jnp.nonzero, static_argnames='size')
>>> nonzero_jit(x, size=3)
(Array([1, 3, 5], dtype=int32),)

If size does not match the true size, the result will be either truncated or padded:

>>> nonzero_jit(x, size=2)  # size < 3: indices are truncated
(Array([1, 3], dtype=int32),)
>>> nonzero_jit(x, size=5)  # size > 3: indices are padded with zeros.
(Array([1, 3, 5, 0, 0], dtype=int32),)

You can specify a custom fill value for the padding using the fill_value argument:

>>> nonzero_jit(x, size=5, fill_value=len(x))
(Array([1, 3, 5, 6, 6], dtype=int32),)
Parameters:
Return type:

tuple[Array, …]

quchip.declarative.qnp.not_equal(x, y, /)

Returns element-wise truth value of x != y.

JAX implementation of numpy.not_equal. This function provides the implementation of the != operator for JAX arrays.

Args:

x: input array or scalar. y: input array or scalar. x and y should either have same shape or be

broadcast compatible.

Returns:

A boolean array containing True where the elements of x != y and False otherwise.

See also:

  • jax.numpy.equal(): Returns element-wise truth value of x == y.

  • jax.numpy.greater_equal(): Returns element-wise truth value of x >= y.

  • jax.numpy.less_equal(): Returns element-wise truth value of x <= y.

  • jax.numpy.greater(): Returns element-wise truth value of x > y.

  • jax.numpy.less(): Returns element-wise truth value of x < y.

Examples:
>>> jnp.not_equal(0., -0.)
Array(False, dtype=bool, weak_type=True)
>>> jnp.not_equal(-2, 2)
Array(True, dtype=bool, weak_type=True)
>>> jnp.not_equal(1, 1.)
Array(False, dtype=bool, weak_type=True)
>>> jnp.not_equal(5, jnp.array(5))
Array(False, dtype=bool, weak_type=True)
>>> x = jnp.array([[1, 2, 3],
...                [4, 5, 6],
...                [7, 8, 9]])
>>> y = jnp.array([1, 5, 9])
>>> jnp.not_equal(x, y)
Array([[False,  True,  True],
       [ True, False,  True],
       [ True,  True, False]], dtype=bool)
>>> x != y
Array([[False,  True,  True],
       [ True, False,  True],
       [ True,  True, False]], dtype=bool)
Parameters:
Return type:

Array

class quchip.declarative.qnp.number

Bases: generic

Abstract base class of all numeric scalar types.

class quchip.declarative.qnp.object_

Bases: generic

Any Python object.

Character code:

'O'

quchip.declarative.qnp.ones(shape, dtype=None, *, device=None)

Create an array full of ones.

JAX implementation of numpy.ones().

Args:

shape: int or sequence of ints specifying the shape of the created array. dtype: optional dtype for the created array; defaults to float32 or float64

depending on the X64 configuration (see default-dtypes).

device: (optional) Device or Sharding

to which the created array will be committed.

Returns:

Array of the specified shape and dtype, on the specified device if specified.

See also:

  • jax.numpy.ones_like()

  • jax.numpy.empty()

  • jax.numpy.zeros()

  • jax.numpy.full()

Examples:
>>> jnp.ones(4)
Array([1., 1., 1., 1.], dtype=float32)
>>> jnp.ones((2, 3), dtype=bool)
Array([[ True,  True,  True],
       [ True,  True,  True]], dtype=bool)
Parameters:

  • shape (Any)

  • dtype (str | type[Any] | dtype | SupportsDType | None)

  • device (Device | Sharding | None)

Return type:

Array

quchip.declarative.qnp.ones_like(a, dtype=None, shape=None, *, device=None)

Create an array of ones with the same shape and dtype as an array.

JAX implementation of numpy.ones_like().

Args:

a: Array-like object with shape and dtype attributes. shape: optionally override the shape of the created array. dtype: optionally override the dtype of the created array. device: (optional) Device or Sharding

to which the created array will be committed.

Returns:

Array of the specified shape and dtype, on the specified device if specified.

See also:

  • jax.numpy.empty()

  • jax.numpy.zeros_like()

  • jax.numpy.ones_like()

  • jax.numpy.full_like()

Examples:
>>> x = jnp.arange(4)
>>> jnp.ones_like(x)
Array([1, 1, 1, 1], dtype=int32)
>>> jnp.ones_like(x, dtype=bool)
Array([ True,  True,  True,  True], dtype=bool)
>>> jnp.ones_like(x, shape=(2, 3))
Array([[1, 1, 1],
       [1, 1, 1]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.outer(a, b, out=None)

Compute the outer product of two arrays.

JAX implementation of numpy.outer().

Args:

a: first input array, if not 1D it will be flattened. b: second input array, if not 1D it will be flattened. out: unsupported by JAX.

Returns:

The outer product of the inputs a and b. Returned array will be of shape (a.size, b.size).

See also:

  • jax.numpy.inner(): compute the inner product of two arrays.

  • jax.numpy.einsum(): Einstein summation.

Examples:
>>> a = jnp.array([1, 2, 3])
>>> b = jnp.array([4, 5, 6])
>>> jnp.outer(a, b)
Array([[ 4,  5,  6],
       [ 8, 10, 12],
       [12, 15, 18]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.packbits(a, axis=None, bitorder='big')

Pack array of bits into a uint8 array.

JAX implementation of numpy.packbits()

Args:

a: N-dimensional array of bits to pack. axis: optional axis along which to pack bits. If not specified, a will

be flattened.

bitorder: "big" (default) or "little": specify whether the bit order

is big-endian or little-endian.

Returns:

A uint8 array of packed values.

See also:

  • jax.numpy.unpackbits(): inverse of packbits.

Examples:

Packing bits in one dimension:

>>> bits = jnp.array([0, 0, 0, 0, 0, 1, 1, 1])
>>> jnp.packbits(bits)
Array([7], dtype=uint8)
>>> 0b00000111  # equivalent bit-wise representation:
7

Optionally specifying little-endian convention:

>>> jnp.packbits(bits, bitorder="little")
Array([224], dtype=uint8)
>>> 0b11100000  # equivalent bit-wise representation
224

If the number of bits is not a multiple of 8, it will be right-padded with zeros:

>>> jnp.packbits(jnp.array([1, 0, 1]))
Array([160], dtype=uint8)
>>> jnp.packbits(jnp.array([1, 0, 1, 0, 0, 0, 0, 0]))
Array([160], dtype=uint8)

For a multi-dimensional input, bits may be packed along a specified axis:

>>> a = jnp.array([[1, 1, 0, 1, 0, 1, 0, 0, 1, 0, 0, 1, 0, 1, 1, 0],
...                [0, 1, 0, 0, 0, 1, 0, 1, 1, 1, 0, 0, 1, 1, 1, 1]])
>>> vals = jnp.packbits(a, axis=1)
>>> vals
Array([[212, 150],
       [ 69, 207]], dtype=uint8)

The inverse of packbits is provided by unpackbits():

>>> jnp.unpackbits(vals, axis=1)
Array([[1, 1, 0, 1, 0, 1, 0, 0, 1, 0, 0, 1, 0, 1, 1, 0],
       [0, 1, 0, 0, 0, 1, 0, 1, 1, 1, 0, 0, 1, 1, 1, 1]], dtype=uint8)
Parameters:
Return type:

Array

quchip.declarative.qnp.pad(array, pad_width, mode='constant', **kwargs)

Add padding to an array.

JAX implementation of numpy.pad().

Args:

array: array to pad. pad_width: specify the pad width for each dimension of an array. Padding widths

may be separately specified for before and after the array. Options are:

  • int or (int,): pad each array dimension with the same number of values both before and after.

  • (before, after): pad each array with before elements before, and after elements after

  • ((before_1, after_1), (before_2, after_2), ... (before_N, after_N)): specify distinct before and after values for each array dimension.

mode: a string or callable. Supported pad modes are:

  • 'constant' (default): pad with a constant value, which defaults to zero.

  • 'empty': pad with empty values (i.e. zero)

  • 'edge': pad with the edge values of the array.

  • 'wrap': pad by wrapping the array.

  • 'linear_ramp': pad with a linear ramp to specified end_values.

  • 'maximum': pad with the maximum value.

  • 'mean': pad with the mean value.

  • 'median': pad with the median value.

  • 'minimum': pad with the minimum value.

  • 'reflect': pad by reflection.

  • 'symmetric': pad by symmetric reflection.

  • <callable>: a callable function. See Notes below.

constant_values: referenced for mode = 'constant'. Specify the constant value

to pad with.

stat_length: referenced for mode in ['maximum', 'mean', 'median', 'minimum'].

An integer or tuple specifying the number of edge values to use when calculating the statistic.

end_values: referenced for mode = 'linear_ramp'. Specify the end values to

ramp the padding values to.

reflect_type: referenced for mode in ['reflect', 'symmetric']. Specify whether

to use even or odd reflection.

Returns:

A padded copy of array.

Notes:

When mode is callable, it should have the following signature:

def pad_func(row: Array, pad_width: tuple[int, int],
             iaxis: int, kwargs: dict) -> Array:
  ...

Here row is a 1D slice of the padded array along axis iaxis, with the pad values filled with zeros. pad_width is a tuple specifying the (before, after) padding sizes, and kwargs are any additional keyword arguments passed to the jax.numpy.pad() function.

Note that while in NumPy, the function should modify row in-place, in JAX the function should return the modified row. In JAX, the custom padding function will be mapped across the padded axis using the jax.vmap() transformation.

See also:

  • jax.numpy.resize(): resize an array

  • jax.numpy.tile(): create a larger array by tiling a smaller array.

  • jax.numpy.repeat(): create a larger array by repeating values of a smaller array.

Examples:

Pad a 1-dimensional array with zeros:

>>> x = jnp.array([10, 20, 30, 40])
>>> jnp.pad(x, 2)
Array([ 0,  0, 10, 20, 30, 40,  0,  0], dtype=int32)
>>> jnp.pad(x, (2, 4))
Array([ 0,  0, 10, 20, 30, 40,  0,  0,  0,  0], dtype=int32)

Pad a 1-dimensional array with specified values:

>>> jnp.pad(x, 2, constant_values=99)
Array([99, 99, 10, 20, 30, 40, 99, 99], dtype=int32)

Pad a 1-dimensional array with the mean array value:

>>> jnp.pad(x, 2, mode='mean')
Array([25, 25, 10, 20, 30, 40, 25, 25], dtype=int32)

Pad a 1-dimensional array with reflected values:

>>> jnp.pad(x, 2, mode='reflect')
Array([30, 20, 10, 20, 30, 40, 30, 20], dtype=int32)

Pad a 2-dimensional array with different paddings in each dimension:

>>> x = jnp.array([[1, 2, 3],
...                [4, 5, 6]])
>>> jnp.pad(x, ((1, 2), (3, 0)))
Array([[0, 0, 0, 0, 0, 0],
       [0, 0, 0, 1, 2, 3],
       [0, 0, 0, 4, 5, 6],
       [0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0]], dtype=int32)

Pad a 1-dimensional array with a custom padding function:

>>> def custom_pad(row, pad_width, iaxis, kwargs):
...   # row represents a 1D slice of the zero-padded array.
...   before, after = pad_width
...   before_value = kwargs.get('before_value', 0)
...   after_value = kwargs.get('after_value', 0)
...   row = row.at[:before].set(before_value)
...   return row.at[len(row) - after:].set(after_value)
>>> x = jnp.array([2, 3, 4])
>>> jnp.pad(x, 2, custom_pad, before_value=-10, after_value=10)
Array([-10, -10,   2,   3,   4,  10,  10], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.partition(a, kth, axis=-1)

Returns a partially-sorted copy of an array.

JAX implementation of numpy.partition(). The JAX version differs from NumPy in the treatment of NaN entries: NaNs which have the negative bit set are sorted to the beginning of the array.

Args:

a: array to be partitioned. kth: static integer index about which to partition the array. axis: static integer axis along which to partition the array; default is -1.

Returns:

A copy of a partitioned at the kth value along axis. The entries before kth are values smaller than take(a, kth, axis), and entries after kth are indices of values larger than take(a, kth, axis)

Note:

The JAX version requires the kth argument to be a static integer rather than a general array. This is implemented via two calls to jax.lax.top_k(). If you’re only accessing the top or bottom k values of the output, it may be more efficient to call jax.lax.top_k() directly.

See Also:

  • jax.numpy.sort(): full sort

  • jax.numpy.argpartition(): indirect partial sort

  • jax.lax.top_k(): directly find the top k entries

  • jax.lax.approx_max_k(): compute the approximate top k entries

  • jax.lax.approx_min_k(): compute the approximate bottom k entries

Examples:
>>> x = jnp.array([6, 8, 4, 3, 1, 9, 7, 5, 2, 3])
>>> kth = 4
>>> x_partitioned = jnp.partition(x, kth)
>>> x_partitioned
Array([1, 2, 3, 3, 4, 9, 8, 7, 6, 5], dtype=int32)

The result is a partially-sorted copy of the input. All values before kth are of smaller than the pivot value, and all values after kth are larger than the pivot value:

>>> smallest_values = x_partitioned[:kth]
>>> pivot_value = x_partitioned[kth]
>>> largest_values = x_partitioned[kth + 1:]
>>> print(smallest_values, pivot_value, largest_values)
[1 2 3 3] 4 [9 8 7 6 5]

Notice that among smallest_values and largest_values, the returned order is arbitrary and implementation-dependent.

Parameters:
Return type:

Array

quchip.declarative.qnp.percentile(a, q, axis=None, out=None, overwrite_input=False, method='linear', keepdims=False, *, interpolation=Deprecated)

Compute the percentile of the data along the specified axis.

JAX implementation of numpy.percentile().

Args:

a: N-dimensional array input. q: scalar or 1-dimensional array specifying the desired quantiles. q

should contain integer or floating point values between 0 and 100.

axis: optional axis or tuple of axes along which to compute the quantile out: not implemented by JAX; will error if not None overwrite_input: not implemented by JAX; will error if not False method: specify the interpolation method to use. Options are one of

["linear", "lower", "higher", "midpoint", "nearest"]. default is linear.

keepdims: if True, then the returned array will have the same number of

dimensions as the input. Default is False.

interpolation: deprecated alias of the method argument. Will result

in a DeprecationWarning if used.

Returns:

An array containing the specified percentiles along the specified axes.

See also:

  • jax.numpy.quantile(): compute the quantile (0.0-1.0)

  • jax.numpy.nanpercentile(): compute the percentile while ignoring NaNs

Examples:

Computing the median and quartiles of a 1D array:

>>> x = jnp.array([0, 1, 2, 3, 4, 5, 6])
>>> q = jnp.array([25, 50, 75])
>>> jnp.percentile(x, q)
Array([1.5, 3. , 4.5], dtype=float32)

Computing the same percentiles with nearest rather than linear interpolation:

>>> jnp.percentile(x, q, method='nearest')
Array([1., 3., 4.], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.permute_dims(a, /, axes)

Permute the axes/dimensions of an array.

JAX implementation of array_api.permute_dims().

Args:

a: input array axes: tuple of integers in range [0, a.ndim) specifying the

axes permutation.

Returns:

a copy of a with axes permuted.

See also:

  • jax.numpy.transpose()

  • jax.numpy.matrix_transpose()

Examples:
>>> a = jnp.array([[1, 2, 3],
...                [4, 5, 6]])
>>> jnp.permute_dims(a, (1, 0))
Array([[1, 4],
       [2, 5],
       [3, 6]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.piecewise(x, condlist, funclist, *args, **kw)

Evaluate a function defined piecewise across the domain.

JAX implementation of numpy.piecewise(), in terms of jax.lax.switch().

Note:

Unlike numpy.piecewise(), jax.numpy.piecewise() requires functions in funclist to be traceable by JAX, as it is implemented via jax.lax.switch().

Args:

x: array of input values. condlist: boolean array or sequence of boolean arrays corresponding to the

functions in funclist. If a sequence of arrays, the length of each array must match the length of x

funclist: list of arrays or functions; must either be the same length as

condlist, or have length len(condlist) + 1, in which case the last entry is the default applied when none of the conditions are True. Alternatively, entries of funclist may be numerical values, in which case they indicate a constant function.

args, kwargs: additional arguments are passed to each function in

funclist.

Returns:

An array which is the result of evaluating the functions on x at the specified conditions.

See also:

  • jax.lax.switch(): choose between N functions based on an index.

  • jax.lax.cond(): choose between two functions based on a boolean condition.

  • jax.numpy.where(): choose between two results based on a boolean mask.

  • jax.lax.select(): choose between two results based on a boolean mask.

  • jax.lax.select_n(): choose between N results based on a boolean mask.

Examples:

Here’s an example of a function which is zero for negative values, and linear for positive values:

>>> x = jnp.array([-4, -3, -2, -1, 0, 1, 2, 3, 4])

>>> condlist = [x < 0, x >= 0]
>>> funclist = [lambda x: 0 * x, lambda x: x]
>>> jnp.piecewise(x, condlist, funclist)
Array([0, 0, 0, 0, 0, 1, 2, 3, 4], dtype=int32)

funclist can also contain a simple scalar value for constant functions:

>>> condlist = [x < 0, x >= 0]
>>> funclist = [0, lambda x: x]
>>> jnp.piecewise(x, condlist, funclist)
Array([0, 0, 0, 0, 0, 1, 2, 3, 4], dtype=int32)

You can specify a default value by appending an extra condition to funclist:

>>> condlist = [x < -1, x > 1]
>>> funclist = [lambda x: 1 + x, lambda x: x - 1, 0]
>>> jnp.piecewise(x, condlist, funclist)
Array([-3, -2,  -1,  0,  0,  0,  1,  2, 3], dtype=int32)

condlist may also be a simple array of scalar conditions, in which case the associated function applies to the whole range

>>> condlist = jnp.array([False, True, False])
>>> funclist = [lambda x: x * 0, lambda x: x * 10, lambda x: x * 100]
>>> jnp.piecewise(x, condlist, funclist)
Array([-40, -30, -20, -10,   0,  10,  20,  30,  40], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.place(arr, mask, vals, *, inplace=True)

Update array elements based on a mask.

JAX implementation of numpy.place().

The semantics of numpy.place() are to modify arrays in-place, which is not possible for JAX’s immutable arrays. The JAX version returns a modified copy of the input, and adds the inplace parameter which must be set to False` by the user as a reminder of this API difference.

Args:

arr: array into which values will be placed. mask: boolean mask with the same size as arr. vals: values to be inserted into arr at the locations indicated

by mask. If too many values are supplied, they will be truncated. If not enough values are supplied, they will be repeated.

inplace: must be set to False to indicate that the input is not modified

in-place, but rather a modified copy is returned.

Returns:

A copy of arr with masked values set to entries from vals.

See Also:

  • jax.numpy.put(): put elements into an array at numerical indices.

  • jax.numpy.ndarray.at(): array updates using NumPy-style indexing

Examples:
>>> x = jnp.zeros((3, 5), dtype=int)
>>> mask = (jnp.arange(x.size) % 3 == 0).reshape(x.shape)
>>> mask
Array([[ True, False, False,  True, False],
       [False,  True, False, False,  True],
       [False, False,  True, False, False]], dtype=bool)

Placing a scalar value:

>>> jnp.place(x, mask, 1, inplace=False)
Array([[1, 0, 0, 1, 0],
       [0, 1, 0, 0, 1],
       [0, 0, 1, 0, 0]], dtype=int32)

In this case, jnp.place is similar to the masked array update syntax:

>>> x.at[mask].set(1)
Array([[1, 0, 0, 1, 0],
       [0, 1, 0, 0, 1],
       [0, 0, 1, 0, 0]], dtype=int32)

place differs when placing values from an array. The array is repeated to fill the masked entries:

>>> vals = jnp.array([1, 3, 5])
>>> jnp.place(x, mask, vals, inplace=False)
Array([[1, 0, 0, 3, 0],
       [0, 5, 0, 0, 1],
       [0, 0, 3, 0, 0]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.poly(seq_of_zeros)

Returns the coefficients of a polynomial for the given sequence of roots.

JAX implementation of numpy.poly().

Args:
seq_of_zeros: A scalar or an array of roots of the polynomial of shape (M,)

or (M, M).

Returns:

An array containing the coefficients of the polynomial. The dtype of the output is always promoted to inexact.

Note:

jax.numpy.poly() differs from numpy.poly():

  • When the input is a scalar, np.poly raises a TypeError, whereas jnp.poly treats scalars the same as length-1 arrays.

  • For complex-valued or square-shaped inputs, jnp.poly always returns complex coefficients, whereas np.poly may return real or complex depending on their values.

See also:

  • jax.numpy.polyfit(): Least squares polynomial fit.

  • jax.numpy.polyval(): Evaluate a polynomial at specific values.

  • jax.numpy.roots(): Computes the roots of a polynomial for given coefficients.

Examples:

Scalar inputs:

>>> jnp.poly(1)
Array([ 1., -1.], dtype=float32)

Input array with integer values:

>>> x = jnp.array([1, 2, 3])
>>> jnp.poly(x)
Array([ 1., -6., 11., -6.], dtype=float32)

Input array with complex conjugates:

>>> x = jnp.array([2, 1+2j, 1-2j])
>>> jnp.poly(x)
Array([  1.+0.j,  -4.+0.j,   9.+0.j, -10.+0.j], dtype=complex64)

Input array as square matrix with real valued inputs:

>>> x = jnp.array([[2, 1, 5],
...                [3, 4, 7],
...                [1, 3, 5]])
>>> jnp.round(jnp.poly(x))
Array([  1.+0.j, -11.-0.j,   9.+0.j, -15.+0.j], dtype=complex64)
Parameters:

seq_of_zeros (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.polyadd(a1, a2)

Returns the sum of the two polynomials.

JAX implementation of numpy.polyadd().

Args:

a1: Array of polynomial coefficients. a2: Array of polynomial coefficients.

Returns:

An array containing the coefficients of the sum of input polynomials.

Note:

jax.numpy.polyadd() only accepts arrays as input unlike numpy.polyadd() which accepts scalar inputs as well.

See also:

  • jax.numpy.polysub(): Computes the difference of two polynomials.

  • jax.numpy.polymul(): Computes the product of two polynomials.

  • jax.numpy.polydiv(): Computes the quotient and remainder of polynomial division.

Examples:
>>> x1 = jnp.array([2, 3])
>>> x2 = jnp.array([5, 4, 1])
>>> jnp.polyadd(x1, x2)
Array([5, 6, 4], dtype=int32)

>>> x3 = jnp.array([[2, 3, 1]])
>>> x4 = jnp.array([[5, 7, 3],
...                 [8, 2, 6]])
>>> jnp.polyadd(x3, x4)
Array([[ 5,  7,  3],
       [10,  5,  7]], dtype=int32)

>>> x5 = jnp.array([1, 3, 5])
>>> x6 = jnp.array([[5, 7, 9],
...                 [8, 6, 4]])
>>> jnp.polyadd(x5, x6)
Traceback (most recent call last):
...
ValueError: Cannot broadcast to shape with fewer dimensions: arr_shape=(2, 3) shape=(2,)
>>> x7 = jnp.array([2])
>>> jnp.polyadd(x6, x7)
Array([[ 5,  7,  9],
       [10,  8,  6]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.polyder(p, m=1)

Returns the coefficients of the derivative of specified order of a polynomial.

JAX implementation of numpy.polyder().

Args:

p: Array of polynomials coefficients. m: Order of differentiation (positive integer). Default is 1. It must be

specified statically.

Returns:

An array of polynomial coefficients representing the derivative.

Note:

jax.numpy.polyder() differs from numpy.polyder() when an integer array is given. NumPy returns the result with dtype int whereas JAX returns the result with dtype float.

See also:

  • jax.numpy.polyint(): Computes the integral of polynomial.

  • jax.numpy.polyval(): Evaluates a polynomial at specific values.

Examples:

The first order derivative of the polynomial \(2 x^3 - 5 x^2 + 3 x - 1\) is \(6 x^2 - 10 x +3\):

>>> p = jnp.array([2, -5, 3, -1])
>>> jnp.polyder(p)
Array([  6., -10.,   3.], dtype=float32)

and its second order derivative is \(12 x - 10\):

>>> jnp.polyder(p, m=2)
Array([ 12., -10.], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.polydiv(u, v, *, trim_leading_zeros=False)

Returns the quotient and remainder of polynomial division.

JAX implementation of numpy.polydiv().

Args:

u: Array of dividend polynomial coefficients. v: Array of divisor polynomial coefficients. trim_leading_zeros: Default is False. If True removes the leading

zeros in the return value to match the result of numpy. But prevents the function from being able to be used in compiled code. Due to differences in accumulation of floating point arithmetic errors, the cutoff for values to be considered zero may lead to inconsistent results between NumPy and JAX, and even between different JAX backends. The result may lead to inconsistent output shapes when trim_leading_zeros=True.

Returns:

A tuple of quotient and remainder arrays. The dtype of the output is always promoted to inexact.

Note:

jax.numpy.polydiv() only accepts arrays as input unlike numpy.polydiv() which accepts scalar inputs as well.

See also:

  • jax.numpy.polyadd(): Computes the sum of two polynomials.

  • jax.numpy.polysub(): Computes the difference of two polynomials.

  • jax.numpy.polymul(): Computes the product of two polynomials.

Examples:
>>> x1 = jnp.array([5, 7, 9])
>>> x2 = jnp.array([4, 1])
>>> np.polydiv(x1, x2)
(array([1.25  , 1.4375]), array([7.5625]))
>>> jnp.polydiv(x1, x2)
(Array([1.25  , 1.4375], dtype=float32), Array([0.    , 0.    , 7.5625], dtype=float32))

If trim_leading_zeros=True, the result matches with np.polydiv’s.

>>> jnp.polydiv(x1, x2, trim_leading_zeros=True)
(Array([1.25  , 1.4375], dtype=float32), Array([7.5625], dtype=float32))
Parameters:
Return type:

tuple[Array, Array]

quchip.declarative.qnp.polyfit(x, y, deg, rcond=None, full=False, w=None, cov=False)

Least squares polynomial fit to data.

Jax implementation of numpy.polyfit().

Given a set of data points (x, y) and degree of polynomial deg, the function finds a polynomial equation of the form:

\[y = p(x) = p[0] x^{deg} + p[1] x^{deg - 1} + ... + p[deg]\]
Args:

x: Array of data points of shape (M,). y: Array of data points of shape (M,) or (M, K). deg: Degree of the polynomials. It must be specified statically. rcond: Relative condition number of the fit. Default value is len(x) * eps.

It must be specified statically.

full: Switch that controls the return value. Default is False which

restricts the return value to the array of polynomial coefficients p. If True, the function returns a tuple (p, resids, rank, s, rcond). It must be specified statically.

w: Array of weights of shape (M,). If None, all data points are considered

to have equal weight. If not None, the weight \(w_i\) is applied to the unsquared residual of \(y_i - \widehat{y}_i\) at \(x_i\), where \(\widehat{y}_i\) is the fitted value of \(y_i\). Default is None.

cov: Boolean or string. If True, returns the covariance matrix scaled

by resids/(M-deg-1) along with polynomial coefficients. If cov='unscaled', returns the unscaled version of covariance matrix. Default is False. cov is ignored if full=True. It must be specified statically.

Returns:

  • An array polynomial coefficients p if full=False and cov=False.

  • A tuple of arrays (p, resids, rank, s, rcond) if full=True. Where

    • p is an array of shape (M,) or (M, K) containing the polynomial coefficients.

    • resids is the sum of squared residual of shape () or (K,).

    • rank is the rank of the matrix x.

    • s is the singular values of the matrix x.

    • rcond as the array.

  • A tuple of arrays (p, C) if full=False and cov=True. Where

    • p is an array of shape (M,) or (M, K) containing the polynomial coefficients.

    • C is the covariance matrix of polynomial coefficients of shape (deg + 1, deg + 1) or (deg + 1, deg + 1, 1).

Note:

Unlike numpy.polyfit() implementation of polyfit, jax.numpy.polyfit() will not warn on rank reduction, which indicates an ill conditioned matrix.

See Also:

  • jax.numpy.poly(): Finds the polynomial coefficients of the given sequence of roots.

  • jax.numpy.polyval(): Evaluate a polynomial at specific values.

  • jax.numpy.roots(): Computes the roots of a polynomial for given coefficients.

Examples:
>>> x = jnp.array([3., 6., 9., 4.])
>>> y = jnp.array([[0, 1, 2],
...                [2, 5, 7],
...                [8, 4, 9],
...                [1, 6, 3]])
>>> p = jnp.polyfit(x, y, 2)
>>> with jnp.printoptions(precision=2, suppress=True):
...   print(p)
[[ 0.2  -0.35 -0.14]
 [-1.17  4.47  2.96]
 [ 1.95 -8.21 -5.93]]

If full=True, returns a tuple of arrays as follows:

>>> p, resids, rank, s, rcond = jnp.polyfit(x, y, 2, full=True)
>>> with jnp.printoptions(precision=2, suppress=True):
...   print("Polynomial Coefficients:", "\n", p, "\n",
...         "Residuals:", resids, "\n",
...         "Rank:", rank, "\n",
...         "s:", s, "\n",
...         "rcond:", rcond)
Polynomial Coefficients:
[[ 0.2  -0.35 -0.14]
[-1.17  4.47  2.96]
[ 1.95 -8.21 -5.93]]
Residuals: [0.37 5.94 0.61]
Rank: 3
s: [1.67 0.47 0.04]
rcond: 4.7683716e-07

If cov=True and full=False, returns a tuple of arrays having polynomial coefficients and covariance matrix.

>>> p, C = jnp.polyfit(x, y, 2, cov=True)
>>> p.shape, C.shape
((3, 3), (3, 3, 3))
Parameters:
Return type:

Array | tuple[Array, …]

quchip.declarative.qnp.polyint(p, m=1, k=None)

Returns the coefficients of the integration of specified order of a polynomial.

JAX implementation of numpy.polyint().

Args:

p: An array of polynomial coefficients. m: Order of integration. Default is 1. It must be specified statically. k: Scalar or array of m integration constant (s).

Returns:

An array of coefficients of integrated polynomial.

See also:

  • jax.numpy.polyder(): Computes the coefficients of the derivative of a polynomial.

  • jax.numpy.polyval(): Evaluates a polynomial at specific values.

Examples:

The first order integration of the polynomial \(12 x^2 + 12 x + 6\) is \(4 x^3 + 6 x^2 + 6 x\).

>>> p = jnp.array([12, 12, 6])
>>> jnp.polyint(p)
Array([4., 6., 6., 0.], dtype=float32)

Since the constant k is not provided, the result included 0 at the end. If the constant k is provided:

>>> jnp.polyint(p, k=4)
Array([4., 6., 6., 4.], dtype=float32)

and the second order integration is \(x^4 + 2 x^3 + 3 x\):

>>> jnp.polyint(p, m=2)
Array([1., 2., 3., 0., 0.], dtype=float32)

When m>=2, the constants k should be provided as an array having m elements. The second order integration of the polynomial \(12 x^2 + 12 x + 6\) with the constants k=[4, 5] is \(x^4 + 2 x^3 + 3 x^2 + 4 x + 5\):

>>> jnp.polyint(p, m=2, k=jnp.array([4, 5]))
Array([1., 2., 3., 4., 5.], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.polymul(a1, a2, *, trim_leading_zeros=False)

Returns the product of two polynomials.

JAX implementation of numpy.polymul().

Args:

a1: 1D array of polynomial coefficients. a2: 1D array of polynomial coefficients. trim_leading_zeros: Default is False. If True removes the leading

zeros in the return value to match the result of numpy. But prevents the function from being able to be used in compiled code. Due to differences in accumulation of floating point arithmetic errors, the cutoff for values to be considered zero may lead to inconsistent results between NumPy and JAX, and even between different JAX backends. The result may lead to inconsistent output shapes when trim_leading_zeros=True.

Returns:

An array of the coefficients of the product of the two polynomials. The dtype of the output is always promoted to inexact.

Note:

jax.numpy.polymul() only accepts arrays as input unlike numpy.polymul() which accepts scalar inputs as well.

See also:

  • jax.numpy.polyadd(): Computes the sum of two polynomials.

  • jax.numpy.polysub(): Computes the difference of two polynomials.

  • jax.numpy.polydiv(): Computes the quotient and remainder of polynomial division.

Examples:
>>> x1 = np.array([2, 1, 0])
>>> x2 = np.array([0, 5, 0, 3])
>>> np.polymul(x1, x2)
array([10,  5,  6,  3,  0])
>>> jnp.polymul(x1, x2)
Array([ 0., 10.,  5.,  6.,  3.,  0.], dtype=float32)

If trim_leading_zeros=True, the result matches with np.polymul’s.

>>> jnp.polymul(x1, x2, trim_leading_zeros=True)
Array([10.,  5.,  6.,  3.,  0.], dtype=float32)

For input arrays of dtype complex:

>>> x3 = np.array([2., 1+2j, 1-2j])
>>> x4 = np.array([0, 5, 0, 3])
>>> np.polymul(x3, x4)
array([10. +0.j,  5.+10.j, 11.-10.j,  3. +6.j,  3. -6.j])
>>> jnp.polymul(x3, x4)
Array([ 0. +0.j, 10. +0.j,  5.+10.j, 11.-10.j,  3. +6.j,  3. -6.j],      dtype=complex64)
>>> jnp.polymul(x3, x4, trim_leading_zeros=True)
Array([10. +0.j,  5.+10.j, 11.-10.j,  3. +6.j,  3. -6.j], dtype=complex64)
Parameters:
Return type:

Array

quchip.declarative.qnp.polysub(a1, a2)

Returns the difference of two polynomials.

JAX implementation of numpy.polysub().

Args:

a1: Array of minuend polynomial coefficients. a2: Array of subtrahend polynomial coefficients.

Returns:

An array containing the coefficients of the difference of two polynomials.

Note:

jax.numpy.polysub() only accepts arrays as input unlike numpy.polysub() which accepts scalar inputs as well.

See also:

  • jax.numpy.polyadd(): Computes the sum of two polynomials.

  • jax.numpy.polymul(): Computes the product of two polynomials.

  • jax.numpy.polydiv(): Computes the quotient and remainder of polynomial division.

Examples:
>>> x1 = jnp.array([2, 3])
>>> x2 = jnp.array([5, 4, 1])
>>> jnp.polysub(x1, x2)
Array([-5, -2,  2], dtype=int32)

>>> x3 = jnp.array([[2, 3, 1]])
>>> x4 = jnp.array([[5, 7, 3],
...                 [8, 2, 6]])
>>> jnp.polysub(x3, x4)
Array([[-5, -7, -3],
       [-6,  1, -5]], dtype=int32)

>>> x5 = jnp.array([1, 3, 5])
>>> x6 = jnp.array([[5, 7, 9],
...                 [8, 6, 4]])
>>> jnp.polysub(x5, x6)
Traceback (most recent call last):
...
ValueError: Cannot broadcast to shape with fewer dimensions: arr_shape=(2, 3) shape=(2,)
>>> x7 = jnp.array([2])
>>> jnp.polysub(x6, x7)
Array([[5, 7, 9],
       [6, 4, 2]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.polyval(p, x, *, unroll=16)

Evaluates the polynomial at specific values.

JAX implementations of numpy.polyval().

For the 1D-polynomial coefficients p of length M, the function returns the value:

\[p_0 x^{M - 1} + p_1 x^{M - 2} + ... + p_{M - 1}\]
Args:

p: An array of polynomial coefficients of shape (M,). x: A number or an array of numbers. unroll: A number used to control the number of unrolled steps with

lax.scan. It must be specified statically.

Returns:

An array of same shape as x.

Note:

The unroll parameter is JAX specific. It does not affect correctness but can have a major impact on performance for evaluating high-order polynomials. The parameter controls the number of unrolled steps with lax.scan inside the jnp.polyval implementation. Consider setting unroll=128 (or even higher) to improve runtime performance on accelerators, at the cost of increased compilation time.

See also:

  • jax.numpy.polyfit(): Least squares polynomial fit.

  • jax.numpy.poly(): Finds the coefficients of a polynomial with given roots.

  • jax.numpy.roots(): Computes the roots of a polynomial for given coefficients.

Examples:
>>> p = jnp.array([2, 5, 1])
>>> jnp.polyval(p, 3)
Array(34., dtype=float32)

If x is a 2D array, polyval returns 2D-array with same shape as that of x:

>>> x = jnp.array([[2, 1, 5],
...                [3, 4, 7],
...                [1, 3, 5]])
>>> jnp.polyval(p, x)
Array([[ 19.,   8.,  76.],
       [ 34.,  53., 134.],
       [  8.,  34.,  76.]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.positive(x, /)

Return element-wise positive values of the input.

JAX implementation of numpy.positive.

Args:

x: input array or scalar

Returns:

An array of same shape and dtype as x containing +x.

Note:

jnp.positive is equivalent to x.copy() and is defined only for the types that support arithmetic operations.

See also:

  • jax.numpy.negative(): Returns element-wise negative values of the input.

  • jax.numpy.sign(): Returns element-wise indication of sign of the input.

Examples:

For real-valued inputs:

>>> x = jnp.array([-5, 4, 7., -9.5])
>>> jnp.positive(x)
Array([-5. ,  4. ,  7. , -9.5], dtype=float32)
>>> x.copy()
Array([-5. ,  4. ,  7. , -9.5], dtype=float32)

For complex inputs:

>>> x1 = jnp.array([1-2j, -3+4j, 5-6j])
>>> jnp.positive(x1)
Array([ 1.-2.j, -3.+4.j,  5.-6.j], dtype=complex64)
>>> x1.copy()
Array([ 1.-2.j, -3.+4.j,  5.-6.j], dtype=complex64)

For uint32:

>>> x2 = jnp.array([6, 0, -4]).astype(jnp.uint32)
>>> x2
Array([         6,          0, 4294967292], dtype=uint32)
>>> jnp.positive(x2)
Array([         6,          0, 4294967292], dtype=uint32)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.pow(x1, x2, /)

Alias of jax.numpy.power()

Parameters:
Return type:

Array

quchip.declarative.qnp.power(x1, x2, /)

Calculate element-wise base x1 exponential of x2.

JAX implementation of numpy.power.

Args:

x1: scalar or array. Specifies the bases. x2: scalar or array. Specifies the exponent. x1 and x2 should either

have same shape or be broadcast compatible.

Returns:

An array containing the base x1 exponentials of x2 with same dtype as input.

Note:

  • When x2 is a concrete integer scalar, jnp.power lowers to jax.lax.integer_pow().

  • When x2 is a traced scalar or an array, jnp.power lowers to jax.lax.pow().

  • jnp.power raises a TypeError for integer type raised to a concrete negative integer power. For a non-concrete power, the operation is invalid and the returned value is implementation-defined.

  • jnp.power returns nan for negative value raised to the power of non-integer values.

See also:

  • jax.lax.pow(): Computes element-wise power, \(x^y\).

  • jax.lax.integer_pow(): Computes element-wise power \(x^y\), where \(y\) is a fixed integer.

  • jax.numpy.float_power(): Computes the first array raised to the power of second array, element-wise, by promoting to the inexact dtype.

  • jax.numpy.pow(): Computes the first array raised to the power of second array, element-wise.

Examples:

Inputs with scalar integers:

>>> jnp.power(4, 3)
Array(64, dtype=int32, weak_type=True)

Inputs with same shape:

>>> x1 = jnp.array([2, 4, 5])
>>> x2 = jnp.array([3, 0.5, 2])
>>> jnp.power(x1, x2)
Array([ 8.,  2., 25.], dtype=float32)

Inputs with broadcast compatibility:

>>> x3 = jnp.array([-2, 3, 1])
>>> x4 = jnp.array([[4, 1, 6],
...                 [1.3, 3, 5]])
>>> jnp.power(x3, x4)
Array([[16.,  3.,  1.],
       [nan, 27.,  1.]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.printoptions(*args, **kwargs)[source]

Alias of numpy.printoptions().

JAX arrays are printed via NumPy, so NumPy’s printoptions configurations will apply to printed JAX arrays.

See the numpy.set_printoptions() documentation for details on the available options and their meanings.

quchip.declarative.qnp.prod(a, axis=None, dtype=None, out=None, keepdims=False, initial=None, where=None, promote_integers=True)

Return product of the array elements over a given axis.

JAX implementation of numpy.prod().

Args:

a: Input array. axis: int or array, default=None. Axis along which the product to be computed.

If None, the product is computed along all the axes.

dtype: The type of the output array. Default=None. keepdims: bool, default=False. If true, reduced axes are left in the result

with size 1.

initial: int or array, Default=None. Initial value for the product. where: int or array, default=None. The elements to be used in the product.

Array should be broadcast compatible to the input.

promote_integersbool, default=True. If True, then integer inputs will be

promoted to the widest available integer dtype, following numpy’s behavior. If False, the result will have the same dtype as the input. promote_integers is ignored if dtype is specified.

out: Unused by JAX.

Returns:

An array of the product along the given axis.

See also:

  • jax.numpy.sum(): Compute the sum of array elements over a given axis.

  • jax.numpy.max(): Compute the maximum of array elements over given axis.

  • jax.numpy.min(): Compute the minimum of array elements over given axis.

Examples:

By default, jnp.prod computes along all the axes.

>>> x = jnp.array([[1, 3, 4, 2],
...                [5, 2, 1, 3],
...                [2, 1, 3, 1]])
>>> jnp.prod(x)
Array(4320, dtype=int32)

If axis=1, product is computed along axis 1.

>>> jnp.prod(x, axis=1)
Array([24, 30,  6], dtype=int32)

If keepdims=True, ndim of the output is equal to that of the input.

>>> jnp.prod(x, axis=1, keepdims=True)
Array([[24],
       [30],
       [ 6]], dtype=int32)

To include only specific elements in the sum, you can use a``where``.

>>> where=jnp.array([[1, 0, 1, 0],
...                  [0, 0, 1, 1],
...                  [1, 1, 1, 0]], dtype=bool)
>>> jnp.prod(x, axis=1, keepdims=True, where=where)
Array([[4],
       [3],
       [6]], dtype=int32)
>>> where = jnp.array([[False],
...                    [False],
...                    [False]])
>>> jnp.prod(x, axis=1, keepdims=True, where=where)
Array([[1],
       [1],
       [1]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.promote_types(a, b)

Returns the type to which a binary operation should cast its arguments.

JAX implementation of numpy.promote_types(). For details of JAX’s type promotion semantics, see type-promotion.

Args:

a: a numpy.dtype or a dtype specifier. b: a numpy.dtype or a dtype specifier.

Returns:

A numpy.dtype object.

Examples:

Type specifiers may be strings, dtypes, or scalar types, and the return value is always a dtype:

>>> jnp.promote_types('int32', 'float32')  # strings
dtype('float32')
>>> jnp.promote_types(jnp.dtype('int32'), jnp.dtype('float32'))  # dtypes
dtype('float32')
>>> jnp.promote_types(jnp.int32, jnp.float32)  # scalar types
dtype('float32')

Built-in scalar types (int, float, or complex) are treated as weakly-typed and will not change the bit width of a strongly-typed counterpart (see discussion in type-promotion):

>>> jnp.promote_types('uint8', int)
dtype('uint8')
>>> jnp.promote_types('float16', float)
dtype('float16')

This differs from the NumPy version of this function, which treats built-in scalar types as equivalent to 64-bit types:

>>> import numpy
>>> numpy.promote_types('uint8', int)
dtype('int64')
>>> numpy.promote_types('float16', float)
dtype('float64')
Parameters:
Return type:

dtype

quchip.declarative.qnp.ptp(a, axis=None, out=None, keepdims=False)

Return the peak-to-peak range along a given axis.

JAX implementation of numpy.ptp().

Args:

a: input array. axis: optional, int or sequence of ints, default=None. Axis along which the

range is computed. If None, the range is computed on the flattened array.

keepdims: bool, default=False. If true, reduced axes are left in the result

with size 1.

out: Unused by JAX.

Returns:

An array with the range of elements along specified axis of input.

Examples:

By default, jnp.ptp computes the range along all axes.

>>> x = jnp.array([[1, 3, 5, 2],
...                [4, 6, 8, 1],
...                [7, 9, 3, 4]])
>>> jnp.ptp(x)
Array(8, dtype=int32)

If axis=1, computes the range along axis 1.

>>> jnp.ptp(x, axis=1)
Array([4, 7, 6], dtype=int32)

To preserve the dimensions of input, you can set keepdims=True.

>>> jnp.ptp(x, axis=1, keepdims=True)
Array([[4],
       [7],
       [6]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.put(a, ind, v, mode=None, *, inplace=True)

Put elements into an array at given indices.

JAX implementation of numpy.put().

The semantics of numpy.put() are to modify arrays in-place, which is not possible for JAX’s immutable arrays. The JAX version returns a modified copy of the input, and adds the inplace parameter which must be set to False` by the user as a reminder of this API difference.

Args:

a: array into which values will be placed. ind: array of indices over the flattened array at which to put values. v: array of values to put into the array. mode: string specifying how to handle out-of-bound indices. Supported values:

  • "clip" (default): clip out-of-bound indices to the final index.

  • "wrap": wrap out-of-bound indices to the beginning of the array.

inplace: must be set to False to indicate that the input is not modified

in-place, but rather a modified copy is returned.

Returns:

A copy of a with specified entries updated.

See Also:

  • jax.numpy.place(): place elements into an array via boolean mask.

  • jax.numpy.ndarray.at(): array updates using NumPy-style indexing.

  • jax.numpy.take(): extract values from an array at given indices.

Examples:
>>> x = jnp.zeros(5, dtype=int)
>>> indices = jnp.array([0, 2, 4])
>>> values = jnp.array([10, 20, 30])
>>> jnp.put(x, indices, values, inplace=False)
Array([10,  0, 20,  0, 30], dtype=int32)

This is equivalent to the following jax.numpy.ndarray.at indexing syntax:

>>> x.at[indices].set(values)
Array([10,  0, 20,  0, 30], dtype=int32)

There are two modes for handling out-of-bound indices. By default they are clipped:

>>> indices = jnp.array([0, 2, 6])
>>> jnp.put(x, indices, values, inplace=False, mode='clip')
Array([10,  0, 20,  0, 30], dtype=int32)

Alternatively, they can be wrapped to the beginning of the array:

>>> jnp.put(x, indices, values, inplace=False, mode='wrap')
Array([10,  30, 20,  0, 0], dtype=int32)

For N-dimensional inputs, the indices refer to the flattened array:

>>> x = jnp.zeros((3, 5), dtype=int)
>>> indices = jnp.array([0, 7, 14])
>>> jnp.put(x, indices, values, inplace=False)
Array([[10,  0,  0,  0,  0],
       [ 0,  0, 20,  0,  0],
       [ 0,  0,  0,  0, 30]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.put_along_axis(arr, indices, values, axis, inplace=True, *, mode=None)

Put values into the destination array by matching 1d index and data slices.

JAX implementation of numpy.put_along_axis().

The semantics of numpy.put_along_axis() are to modify arrays in-place, which is not possible for JAX’s immutable arrays. The JAX version returns a modified copy of the input, and adds the inplace parameter which must be set to False` by the user as a reminder of this API difference.

Args:

arr: array into which values will be put. indices: array of indices at which to put values. values: array of values to put into the array. axis: the axis along which to put values. If not specified, the array will

be flattened before indexing is applied.

inplace: must be set to False to indicate that the input is not modified

in-place, but rather a modified copy is returned.

mode: Out-of-bounds indexing mode. For more discussion of mode options,

see jax.numpy.ndarray.at.

Returns:

A copy of a with specified entries updated.

See Also:

  • jax.numpy.put(): put elements into an array at given indices.

  • jax.numpy.place(): place elements into an array via boolean mask.

  • jax.numpy.ndarray.at(): array updates using NumPy-style indexing.

  • jax.numpy.take(): extract values from an array at given indices.

  • jax.numpy.take_along_axis(): extract values from an array along an axis.

Examples:
>>> from jax import numpy as jnp
>>> a = jnp.array([[10, 30, 20], [60, 40, 50]])
>>> i = jnp.argmax(a, axis=1, keepdims=True)
>>> print(i)
[[1]
 [0]]
>>> b = jnp.put_along_axis(a, i, 99, axis=1, inplace=False)
>>> print(b)
[[10 99 20]
 [99 40 50]]
Parameters:
Return type:

Array

quchip.declarative.qnp.quantile(a, q, axis=None, out=None, overwrite_input=False, method='linear', keepdims=False, *, interpolation=Deprecated)

Compute the quantile of the data along the specified axis.

JAX implementation of numpy.quantile().

Args:

a: N-dimensional array input. q: scalar or 1-dimensional array specifying the desired quantiles. q

should contain floating-point values between 0.0 and 1.0.

axis: optional axis or tuple of axes along which to compute the quantile out: not implemented by JAX; will error if not None overwrite_input: not implemented by JAX; will error if not False method: specify the interpolation method to use. Options are one of

["linear", "lower", "higher", "midpoint", "nearest"]. default is linear.

keepdims: if True, then the returned array will have the same number of

dimensions as the input. Default is False.

interpolation: deprecated alias of the method argument. Will result

in a DeprecationWarning if used.

Returns:

An array containing the specified quantiles along the specified axes.

See also:

  • jax.numpy.nanquantile(): compute the quantile while ignoring NaNs

  • jax.numpy.percentile(): compute the percentile (0-100)

Examples:

Computing the median and quartiles of an array, with linear interpolation:

>>> x = jnp.arange(10)
>>> q = jnp.array([0.25, 0.5, 0.75])
>>> jnp.quantile(x, q)
Array([2.25, 4.5 , 6.75], dtype=float32)

Computing the quartiles using nearest-value interpolation:

>>> jnp.quantile(x, q, method='nearest')
Array([2., 4., 7.], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.rad2deg(x, /)

Convert angles from radians to degrees.

JAX implementation of numpy.rad2deg.

The angle in radians is converted to degrees by:

\[rad2deg(x) = x * \frac{180}{pi}\]
Args:

x: scalar or array. Specifies the angle in radians.

Returns:

An array containing the angles in degrees.

See also:

  • jax.numpy.deg2rad() and jax.numpy.radians(): Converts the angles from degrees to radians.

  • jax.numpy.degrees(): Alias of rad2deg.

Examples:
>>> pi = jnp.pi
>>> x = jnp.array([pi/4, pi/2, 2*pi/3])
>>> jnp.rad2deg(x)
Array([ 45.     ,  90.     , 120.00001], dtype=float32)
>>> x * 180 / pi
Array([ 45.,  90., 120.], dtype=float32)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.radians(x, /)

Alias of jax.numpy.deg2rad()

Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.ravel(a, order='C', *, out_sharding=None)

Flatten array into a 1-dimensional shape.

JAX implementation of numpy.ravel(), implemented in terms of jax.lax.reshape().

ravel(arr, order=order) is equivalent to reshape(arr, -1, order=order).

Args:

a: array to be flattened. order: 'F' or 'C', specifies whether the reshape should apply column-major

(fortran-style, "F") or row-major (C-style, "C") order; default is "C". JAX does not support order=”A” or order=”K”.

Returns:

flattened copy of input array.

Notes:

Unlike numpy.ravel(), jax.numpy.ravel() will return a copy rather than a view of the input array. However, under JIT, the compiler will optimize-away such copies when possible, so this doesn’t have performance impacts in practice.

See Also:

  • jax.Array.ravel(): equivalent functionality via an array method.

  • jax.numpy.reshape(): general array reshape.

Examples:
>>> x = jnp.array([[1, 2, 3],
...                [4, 5, 6]])

By default, ravel in C-style, row-major order

>>> jnp.ravel(x)
Array([1, 2, 3, 4, 5, 6], dtype=int32)

Optionally ravel in Fortran-style, column-major:

>>> jnp.ravel(x, order='F')
Array([1, 4, 2, 5, 3, 6], dtype=int32)

For convenience, the same functionality is available via the jax.Array.ravel() method:

>>> x.ravel()
Array([1, 2, 3, 4, 5, 6], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.ravel_multi_index(multi_index, dims, mode='raise', order='C')

Convert multi-dimensional indices into flat indices.

JAX implementation of numpy.ravel_multi_index()

Args:

multi_index: sequence of integer arrays containing indices in each dimension. dims: sequence of integer sizes; must have len(dims) == len(multi_index) mode: how to handle out-of bound indices. Options are

  • "raise" (default): raise a ValueError. This mode is incompatible with jit() or other JAX transformations.

  • "clip": clip out-of-bound indices to valid range.

  • "wrap": wrap out-of-bound indices to valid range.

order: "C" (default) or "F", specify whether to assume C-style

row-major order or Fortran-style column-major order.

Returns:

array of flattened indices

See also:

jax.numpy.unravel_index(): inverse of this function.

Examples:

Define a 2-dimensional array and a sequence of indices of even values:

>>> x = jnp.array([[2., 3., 4.],
...                [5., 6., 7.]])
>>> indices = jnp.where(x % 2 == 0)
>>> indices
(Array([0, 0, 1], dtype=int32), Array([0, 2, 1], dtype=int32))
>>> x[indices]
Array([2., 4., 6.], dtype=float32)

Compute the flattened indices:

>>> indices_flat = jnp.ravel_multi_index(indices, x.shape)
>>> indices_flat
Array([0, 2, 4], dtype=int32)

These flattened indices can be used to extract the same values from the flattened x array:

>>> x_flat = x.ravel()
>>> x_flat
Array([2., 3., 4., 5., 6., 7.], dtype=float32)
>>> x_flat[indices_flat]
Array([2., 4., 6.], dtype=float32)

The original indices can be recovered with unravel_index():

>>> jnp.unravel_index(indices_flat, x.shape)
(Array([0, 0, 1], dtype=int32), Array([0, 2, 1], dtype=int32))
Parameters:
Return type:

Array

quchip.declarative.qnp.real(val, /)

Return element-wise real part of the complex argument.

JAX implementation of numpy.real.

Args:

val: input array or scalar.

Returns:

An array containing the real part of the elements of val.

See also:

  • jax.numpy.conjugate() and jax.numpy.conj(): Returns the element-wise complex-conjugate of the input.

  • jax.numpy.imag(): Returns the element-wise imaginary part of the complex argument.

Examples:
>>> jnp.real(5)
Array(5, dtype=int32, weak_type=True)
>>> jnp.real(2j)
Array(0., dtype=float32, weak_type=True)
>>> x = jnp.array([3-2j, 4+7j, -2j])
>>> jnp.real(x)
Array([ 3.,  4., -0.], dtype=float32)
Parameters:

val (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.reciprocal(x, /)

Calculate element-wise reciprocal of the input.

JAX implementation of numpy.reciprocal.

The reciprocal is calculated by 1/x.

Args:

x: input array or scalar.

Returns:

An array of same shape as x containing the reciprocal of each element of x.

Note:

For integer inputs, np.reciprocal returns rounded integer output, while jnp.reciprocal promotes integer inputs to floating point.

Examples:
>>> jnp.reciprocal(2)
Array(0.5, dtype=float32, weak_type=True)
>>> jnp.reciprocal(0.)
Array(inf, dtype=float32, weak_type=True)
>>> x = jnp.array([1, 5., 4.])
>>> jnp.reciprocal(x)
Array([1.  , 0.2 , 0.25], dtype=float32)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.remainder(x1, x2, /)

Returns element-wise remainder of the division.

JAX implementation of numpy.remainder.

Args:

x1: scalar or array. Specifies the dividend. x2: scalar or array. Specifies the divisor. x1 and x2 should either

have same shape or be broadcast compatible.

Returns:

An array containing the remainder of element-wise division of x1 by x2 with same sign as the elements of x2.

Note:

The result of jnp.remainder is equivalent to x1 - x2 * jnp.floor(x1 / x2).

See also:

  • jax.numpy.mod(): Returns the element-wise remainder of the division.

  • jax.numpy.fmod(): Calculates the element-wise floating-point modulo operation.

  • jax.numpy.divmod(): Calculates the integer quotient and remainder of x1 by x2, element-wise.

Examples:
>>> x1 = jnp.array([[3, -1, 4],
...                 [8, 5, -2]])
>>> x2 = jnp.array([2, 3, -5])
>>> jnp.remainder(x1, x2)
Array([[ 1,  2, -1],
       [ 0,  2, -2]], dtype=int32)
>>> x1 - x2 * jnp.floor(x1 / x2)
Array([[ 1.,  2., -1.],
       [ 0.,  2., -2.]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.repeat(a, repeats, axis=None, *, total_repeat_length=None, out_sharding=None)

Construct an array from repeated elements.

JAX implementation of numpy.repeat().

Args:

a: N-dimensional array repeats: 1D integer array specifying the number of repeats. Must match the

length of the repeated axis.

axis: integer specifying the axis of a along which to construct the

repeated array. If None (default) then a is first flattened.

total_repeat_length: this must be specified statically for jnp.repeat

to be compatible with jit() and other JAX transformations. If sum(repeats) is larger than the specified total_repeat_length, the remaining values will be discarded. If sum(repeats) is smaller than total_repeat_length, the final value will be repeated.

Returns:

an array constructed from repeated values of a.

See Also:

  • jax.numpy.tile(): repeat a full array rather than individual values.

Examples:

Repeat each value twice along the last axis:

>>> a = jnp.array([[1, 2],
...                [3, 4]])
>>> jnp.repeat(a, 2, axis=-1)
Array([[1, 1, 2, 2],
       [3, 3, 4, 4]], dtype=int32)

If axis is not specified, the input array will be flattened:

>>> jnp.repeat(a, 2)
Array([1, 1, 2, 2, 3, 3, 4, 4], dtype=int32)

Pass an array to repeats to repeat each value a different number of times:

>>> repeats = jnp.array([2, 3])
>>> jnp.repeat(a, repeats, axis=1)
Array([[1, 1, 2, 2, 2],
       [3, 3, 4, 4, 4]], dtype=int32)

In order to use repeat within jit and other JAX transformations, the size of the output must be specified statically using total_repeat_length:

>>> jit_repeat = jax.jit(jnp.repeat, static_argnames=['axis', 'total_repeat_length'])
>>> jit_repeat(a, repeats, axis=1, total_repeat_length=5)
Array([[1, 1, 2, 2, 2],
       [3, 3, 4, 4, 4]], dtype=int32)

If total_repeat_length is smaller than sum(repeats), the result will be truncated:

>>> jit_repeat(a, repeats, axis=1, total_repeat_length=4)
Array([[1, 1, 2, 2],
       [3, 3, 4, 4]], dtype=int32)

If it is larger, then the additional entries will be filled with the final value:

>>> jit_repeat(a, repeats, axis=1, total_repeat_length=7)
Array([[1, 1, 2, 2, 2, 2, 2],
       [3, 3, 4, 4, 4, 4, 4]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.reshape(a, shape, order='C', *, copy=None, out_sharding=None)

Return a reshaped copy of an array.

JAX implementation of numpy.reshape(), implemented in terms of jax.lax.reshape().

Args:

a: input array to reshape shape: integer or sequence of integers giving the new shape, which must match the

size of the input array. If any single dimension is given size -1, it will be replaced with a value such that the output has the correct size.

order: 'F' or 'C', specifies whether the reshape should apply column-major

(fortran-style, "F") or row-major (C-style, "C") order; default is "C". JAX does not support order="A".

copy: unused by JAX; JAX always returns a copy, though under JIT the compiler

may optimize such copies away.

Returns:

reshaped copy of input array with the specified shape.

Notes:

Unlike numpy.reshape(), jax.numpy.reshape() will return a copy rather than a view of the input array. However, under JIT, the compiler will optimize-away such copies when possible, so this doesn’t have performance impacts in practice.

See Also:

  • jax.Array.reshape(): equivalent functionality via an array method.

  • jax.numpy.ravel(): flatten an array into a 1D shape.

  • jax.numpy.squeeze(): remove one or more length-1 axes from an array’s shape.

Examples:
>>> x = jnp.array([[1, 2, 3],
...                [4, 5, 6]])
>>> jnp.reshape(x, 6)
Array([1, 2, 3, 4, 5, 6], dtype=int32)
>>> jnp.reshape(x, (3, 2))
Array([[1, 2],
       [3, 4],
       [5, 6]], dtype=int32)

You can use -1 to automatically compute a shape that is consistent with the input size:

>>> jnp.reshape(x, -1)  # -1 is inferred to be 6
Array([1, 2, 3, 4, 5, 6], dtype=int32)
>>> jnp.reshape(x, (-1, 2))  # -1 is inferred to be 3
Array([[1, 2],
       [3, 4],
       [5, 6]], dtype=int32)

The default ordering of axes in the reshape is C-style row-major ordering. To use Fortran-style column-major ordering, specify order='F':

>>> jnp.reshape(x, 6, order='F')
Array([1, 4, 2, 5, 3, 6], dtype=int32)
>>> jnp.reshape(x, (3, 2), order='F')
Array([[1, 5],
       [4, 3],
       [2, 6]], dtype=int32)

For convenience, this functionality is also available via the jax.Array.reshape() method:

>>> x.reshape(3, 2)
Array([[1, 2],
       [3, 4],
       [5, 6]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.resize(a, new_shape)

Return a new array with specified shape.

JAX implementation of numpy.resize().

Args:

a: input array or scalar. new_shape: int or tuple of ints. Specifies the shape of the resized array.

Returns:

A resized array with specified shape. The elements of a are repeated in the resized array, if the resized array is larger than the original array.

See also:

  • jax.numpy.reshape(): Returns a reshaped copy of an array.

  • jax.numpy.repeat(): Constructs an array from repeated elements.

Examples:
>>> x = jnp.array([1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> jnp.resize(x, (3, 3))
Array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]], dtype=int32)
>>> jnp.resize(x, (3, 4))
Array([[1, 2, 3, 4],
       [5, 6, 7, 8],
       [9, 1, 2, 3]], dtype=int32)
>>> jnp.resize(4, (3, 2))
Array([[4, 4],
       [4, 4],
       [4, 4]], dtype=int32, weak_type=True)
Parameters:
Return type:

Array

quchip.declarative.qnp.result_type(*args)

Return the result of applying JAX promotion rules to the inputs.

JAX implementation of numpy.result_type().

JAX’s dtype promotion behavior is described in type-promotion.

Args:

args: one or more arrays or dtype-like objects.

Returns:

A numpy.dtype instance representing the result of type promotion for the inputs.

Examples:

Inputs can be dtype specifiers:

>>> jnp.result_type('int32', 'float32')
dtype('float32')
>>> jnp.result_type(np.uint16, np.dtype('int32'))
dtype('int32')

Inputs may also be scalars or arrays:

>>> jnp.result_type(1.0, jnp.bfloat16(2))
dtype(bfloat16)
>>> jnp.result_type(jnp.arange(4), jnp.zeros(4))
dtype('float32')

Be aware that the result type will be canonicalized based on the state of the jax_enable_x64 configuration flag, meaning that 64-bit types may be downcast to 32-bit:

>>> jnp.result_type('float64')
dtype('float32')

For details on 64-bit values, refer to Sharp bits - double precision:

Parameters:

args (Any)

Return type:

dtype

quchip.declarative.qnp.right_shift(x1, x2, /)

Right shift the bits of x1 to the amount specified in x2.

JAX implementation of numpy.right_shift.

Args:

x1: Input array, only accepts unsigned integer subtypes x2: The amount of bits to shift each element in x1 to the right, only accepts

integer subtypes

Returns:

An array-like object containing the right shifted elements of x1 by the amount specified in x2, with the same shape as the broadcasted shape of x1 and x2.

Note:

If x1.shape != x2.shape, they must be compatible for broadcasting to a shared shape, this shared shape will also be the shape of the output. Right shifting a scalar x1 by scalar x2 is equivalent to x1 // 2**x2.

Examples:
>>> def print_binary(x):
...   return [bin(int(val)) for val in x]

>>> x1 = jnp.array([1, 2, 4, 8])
>>> print_binary(x1)
['0b1', '0b10', '0b100', '0b1000']
>>> x2 = 1
>>> result = jnp.right_shift(x1, x2)
>>> result
Array([0, 1, 2, 4], dtype=int32)
>>> print_binary(result)
['0b0', '0b1', '0b10', '0b100']

>>> x1 = 16
>>> print_binary([x1])
['0b10000']
>>> x2 = jnp.array([1, 2, 3, 4])
>>> result = jnp.right_shift(x1, x2)
>>> result
Array([8, 4, 2, 1], dtype=int32)
>>> print_binary(result)
['0b1000', '0b100', '0b10', '0b1']
Parameters:
Return type:

Array

quchip.declarative.qnp.rint(x, /)

Rounds the elements of x to the nearest integer

JAX implementation of numpy.rint.

Args:

x: Input array

Returns:

An array-like object containing the rounded elements of x. Always promotes to inexact.

Note:

If an element of x is exactly half way, e.g. 0.5 or 1.5, rint will round to the nearest even integer.

Examples:
>>> x1 = jnp.array([5, 4, 7])
>>> jnp.rint(x1)
Array([5., 4., 7.], dtype=float32)

>>> x2 = jnp.array([-2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5, 4.5])
>>> jnp.rint(x2)
Array([-2., -2., -0.,  0.,  2.,  2.,  4.,  4.], dtype=float32)

>>> x3 = jnp.array([-2.5+3.5j, 4.5-0.5j])
>>> jnp.rint(x3)
Array([-2.+4.j,  4.-0.j], dtype=complex64)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.roll(a, shift, axis=None)

Roll the elements of an array along a specified axis.

JAX implementation of numpy.roll().

Args:

a: input array. shift: the number of positions to shift the specified axis. If an integer,

all axes are shifted by the same amount. If a tuple, the shift for each axis is specified individually.

axis: the axis or axes to roll. If None, the array is flattened, shifted,

and then reshaped to its original shape.

Returns:

A copy of a with elements rolled along the specified axis or axes.

See also:

  • jax.numpy.rollaxis(): roll the specified axis to a given position.

Examples:
>>> a = jnp.array([0, 1, 2, 3, 4, 5])
>>> jnp.roll(a, 2)
Array([4, 5, 0, 1, 2, 3], dtype=int32)

Roll elements along a specific axis:

>>> a = jnp.array([[ 0,  1,  2,  3],
...                [ 4,  5,  6,  7],
...                [ 8,  9, 10, 11]])
>>> jnp.roll(a, 1, axis=0)
Array([[ 8,  9, 10, 11],
       [ 0,  1,  2,  3],
       [ 4,  5,  6,  7]], dtype=int32)
>>> jnp.roll(a, [2, 3], axis=[0, 1])
Array([[ 5,  6,  7,  4],
       [ 9, 10, 11,  8],
       [ 1,  2,  3,  0]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.rollaxis(a, axis, start=0)

Roll the specified axis to a given position.

JAX implementation of numpy.rollaxis().

This function exists for compatibility with NumPy, but in most cases the newer jax.numpy.moveaxis() instead, because the meaning of its arguments is more intuitive.

Args:

a: input array. axis: index of the axis to roll forward. start: index toward which the axis will be rolled (default = 0). After

normalizing negative axes, if start <= axis, the axis is rolled to the start index; if start > axis, the axis is rolled until the position before start.

Returns:

Copy of a with rolled axis.

Notes:

Unlike numpy.rollaxis(), jax.numpy.rollaxis() will return a copy rather than a view of the input array. However, under JIT, the compiler will optimize away such copies when possible, so this doesn’t have performance impacts in practice.

See also:

  • jax.numpy.moveaxis(): newer API with clearer semantics than rollaxis; this should be preferred to rollaxis in most cases.

  • jax.numpy.swapaxes(): swap two axes.

  • jax.numpy.transpose(): general permutation of axes.

Examples:
>>> a = jnp.ones((2, 3, 4, 5))

Roll axis 2 to the start of the array:

>>> jnp.rollaxis(a, 2).shape
(4, 2, 3, 5)

Roll axis 1 to the end of the array:

>>> jnp.rollaxis(a, 1, a.ndim).shape
(2, 4, 5, 3)

Equivalent of these two with moveaxis()

>>> jnp.moveaxis(a, 2, 0).shape
(4, 2, 3, 5)
>>> jnp.moveaxis(a, 1, -1).shape
(2, 4, 5, 3)
Parameters:
Return type:

Array

quchip.declarative.qnp.roots(p, *, strip_zeros=True)

Returns the roots of a polynomial given the coefficients p.

JAX implementations of numpy.roots().

Args:

p: Array of polynomial coefficients having rank-1. strip_zeros : bool, default=True. If True, then leading zeros in the

coefficients will be stripped, similar to numpy.roots(). If set to False, leading zeros will not be stripped, and undefined roots will be represented by NaN values in the function output. strip_zeros must be set to False for the function to be compatible with jax.jit() and other JAX transformations.

Returns:

An array containing the roots of the polynomial.

Note:

Unlike np.roots of this function, the jnp.roots returns the roots in a complex array regardless of the values of the roots.

See Also:

  • jax.numpy.poly(): Finds the polynomial coefficients of the given sequence of roots.

  • jax.numpy.polyfit(): Least squares polynomial fit to data.

  • jax.numpy.polyval(): Evaluate a polynomial at specific values.

Examples:
>>> coeffs = jnp.array([0, 1, 2])

The default behavior matches numpy and strips leading zeros:

>>> jnp.roots(coeffs)
Array([-2.+0.j], dtype=complex64)

With strip_zeros=False, extra roots are set to NaN:

>>> jnp.roots(coeffs, strip_zeros=False)
Array([-2. +0.j, nan+nanj], dtype=complex64)
Parameters:
Return type:

Array

quchip.declarative.qnp.rot90(m, k=1, axes=(0, 1))

Rotate an array by 90 degrees counterclockwise in the plane specified by axes.

JAX implementation of numpy.rot90().

Args:

m: input array. Must have m.ndim >= 2. k: int, optional, default=1. Specifies the number of times the array is rotated.

For negative values of k, the array is rotated in clockwise direction.

axes: tuple of 2 integers, optional, default= (0, 1). The axes define the plane

in which the array is rotated. Both the axes must be different.

Returns:

An array containing the copy of the input, m rotated by 90 degrees.

See also:

  • jax.numpy.flip(): reverse the order along the given axis

  • jax.numpy.fliplr(): reverse the order along axis 1 (left/right)

  • jax.numpy.flipud(): reverse the order along axis 0 (up/down)

Examples:
>>> m = jnp.array([[1, 2, 3],
...                [4, 5, 6]])
>>> jnp.rot90(m)
Array([[3, 6],
       [2, 5],
       [1, 4]], dtype=int32)
>>> jnp.rot90(m, k=2)
Array([[6, 5, 4],
       [3, 2, 1]], dtype=int32)

jnp.rot90(m, k=1, axes=(1, 0)) is equivalent to jnp.rot90(m, k=-1, axes(0,1)).

>>> jnp.rot90(m, axes=(1, 0))
Array([[4, 1],
       [5, 2],
       [6, 3]], dtype=int32)
>>> jnp.rot90(m, k=-1, axes=(0, 1))
Array([[4, 1],
       [5, 2],
       [6, 3]], dtype=int32)

when input array has ndim>2:

>>> m1 = jnp.array([[[1, 2, 3],
...                  [4, 5, 6]],
...                 [[7, 8, 9],
...                  [10, 11, 12]]])
>>> jnp.rot90(m1, k=1, axes=(2, 1))
Array([[[ 4,  1],
        [ 5,  2],
        [ 6,  3]],

       [[10,  7],
        [11,  8],
        [12,  9]]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.round(a, decimals=0, out=None)

Round input evenly to the given number of decimals.

JAX implementation of numpy.round().

Args:

a: input array or scalar. decimals: int, default=0. Number of decimal points to which the input needs

to be rounded. It must be specified statically. Not implemented for decimals < 0.

out: Unused by JAX.

Returns:

An array containing the rounded values to the specified decimals with same shape and dtype as a.

Note:

jnp.round rounds to the nearest even integer for the values exactly halfway between rounded decimal values.

See also:

  • jax.numpy.floor(): Rounds the input to the nearest integer downwards.

  • jax.numpy.ceil(): Rounds the input to the nearest integer upwards.

  • jax.numpy.fix() and :func:numpy.trunc`: Rounds the input to the nearest integer towards zero.

Examples:
>>> x = jnp.array([1.532, 3.267, 6.149])
>>> jnp.round(x)
Array([2., 3., 6.], dtype=float32)
>>> jnp.round(x, decimals=2)
Array([1.53, 3.27, 6.15], dtype=float32)

For values exactly halfway between rounded values:

>>> x1 = jnp.array([10.5, 21.5, 12.5, 31.5])
>>> jnp.round(x1)
Array([10., 22., 12., 32.], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.save(file, arr, allow_pickle=True, fix_imports=<no value>)

Save an array to a binary file in NumPy .npy format.

Parameters:

  • file (file, str, or pathlib.Path) – File or filename to which the data is saved. If file is a file-object, then the filename is unchanged. If file is a string or Path, a .npy extension will be appended to the filename if it does not already have one.

  • arr (array_like) – Array data to be saved.

  • allow_pickle (bool, optional) – Allow saving object arrays using Python pickles. Reasons for disallowing pickles include security (loading pickled data can execute arbitrary code) and portability (pickled objects may not be loadable on different Python installations, for example if the stored objects require libraries that are not available, and not all pickled data is compatible between different versions of Python). Default: True

  • fix_imports (bool, optional) –

    The fix_imports flag is deprecated and has no effect.

    Deprecated since version 2.1: This flag is ignored since NumPy 1.17 and was only needed to support loading some files in Python 2 written in Python 3.

See also

savez

Save several arrays into a .npz archive

savetxt, load

Notes

For a description of the .npy format, see numpy.lib.format.

Any data saved to the file is appended to the end of the file.

Examples

>>> import numpy as np

>>> from tempfile import TemporaryFile
>>> outfile = TemporaryFile()

>>> x = np.arange(10)
>>> np.save(outfile, x)

>>> _ = outfile.seek(0) # Only needed to simulate closing & reopening file
>>> np.load(outfile)
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])

>>> with open('test.npy', 'wb') as f:
...     np.save(f, np.array([1, 2]))
...     np.save(f, np.array([1, 3]))
>>> with open('test.npy', 'rb') as f:
...     a = np.load(f)
...     b = np.load(f)
>>> print(a, b)
# [1 2] [1 3]
quchip.declarative.qnp.savez(file, *args, allow_pickle=True, **kwds)

Save several arrays into a single file in uncompressed .npz format.

Provide arrays as keyword arguments to store them under the corresponding name in the output file: savez(fn, x=x, y=y).

If arrays are specified as positional arguments, i.e., savez(fn, x, y), their names will be arr_0, arr_1, etc.

Parameters:

  • file (file, str, or pathlib.Path) – Either the filename (string) or an open file (file-like object) where the data will be saved. If file is a string or a Path, the .npz extension will be appended to the filename if it is not already there.

  • args (Arguments, optional) – Arrays to save to the file. Please use keyword arguments (see kwds below) to assign names to arrays. Arrays specified as args will be named “arr_0”, “arr_1”, and so on.

  • allow_pickle (bool, optional) – Allow saving object arrays using Python pickles. Reasons for disallowing pickles include security (loading pickled data can execute arbitrary code) and portability (pickled objects may not be loadable on different Python installations, for example if the stored objects require libraries that are not available, and not all pickled data is compatible between different versions of Python). Default: True

  • kwds (Keyword arguments, optional) – Arrays to save to the file. Each array will be saved to the output file with its corresponding keyword name.

Return type:

None

See also

save

Save a single array to a binary file in NumPy format.

savetxt

Save an array to a file as plain text.

savez_compressed

Save several arrays into a compressed .npz archive

Notes

The .npz file format is a zipped archive of files named after the variables they contain. The archive is not compressed and each file in the archive contains one variable in .npy format. For a description of the .npy format, see numpy.lib.format.

When opening the saved .npz file with load a ~lib.npyio.NpzFile object is returned. This is a dictionary-like object which can be queried for its list of arrays (with the .files attribute), and for the arrays themselves.

Keys passed in kwds are used as filenames inside the ZIP archive. Therefore, keys should be valid filenames; e.g., avoid keys that begin with / or contain ..

When naming variables with keyword arguments, it is not possible to name a variable file, as this would cause the file argument to be defined twice in the call to savez.

Examples

>>> import numpy as np
>>> from tempfile import TemporaryFile
>>> outfile = TemporaryFile()
>>> x = np.arange(10)
>>> y = np.sin(x)

Using savez with *args, the arrays are saved with default names.

>>> np.savez(outfile, x, y)
>>> _ = outfile.seek(0) # Only needed to simulate closing & reopening file
>>> npzfile = np.load(outfile)
>>> npzfile.files
['arr_0', 'arr_1']
>>> npzfile['arr_0']
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])

Using savez with **kwds, the arrays are saved with the keyword names.

>>> outfile = TemporaryFile()
>>> np.savez(outfile, x=x, y=y)
>>> _ = outfile.seek(0)
>>> npzfile = np.load(outfile)
>>> sorted(npzfile.files)
['x', 'y']
>>> npzfile['x']
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
quchip.declarative.qnp.searchsorted(a, v, side='left', sorter=None, *, method='scan')

Perform a binary search within a sorted array.

JAX implementation of numpy.searchsorted().

This will return the indices within a sorted array a where values in v can be inserted to maintain its sort order.

Args:

a: one-dimensional array, assumed to be in sorted order unless sorter is specified. v: N-dimensional array of query values side: 'left' (default) or 'right'; specifies whether insertion indices will be

to the left or the right in case of ties.

sorter: optional array of indices specifying the sort order of a. If specified,

then the algorithm assumes that a[sorter] is in sorted order.

method: one of 'scan' (default), 'scan_unrolled', 'sort' or 'compare_all'.

See Note below.

Returns:

Array of insertion indices of shape v.shape.

Note:

The method argument controls the algorithm used to compute the insertion indices.

  • 'scan' (the default) tends to be more performant on CPU, particularly when a is very large.

  • 'scan_unrolled' is more performant on GPU at the expense of additional compile time.

  • 'sort' is often more performant on accelerator backends like GPU and TPU, particularly when v is very large.

  • 'compare_all' tends to be the most performant when a is very small.

Examples:

Searching for a single value:

>>> a = jnp.array([1, 2, 2, 3, 4, 5, 5])
>>> jnp.searchsorted(a, 2)
Array(1, dtype=int32)
>>> jnp.searchsorted(a, 2, side='right')
Array(3, dtype=int32)

Searching for a batch of values:

>>> vals = jnp.array([0, 3, 8, 1.5, 2])
>>> jnp.searchsorted(a, vals)
Array([0, 3, 7, 1, 1], dtype=int32)

Optionally, the sorter argument can be used to find insertion indices into an array sorted via jax.numpy.argsort():

>>> a = jnp.array([4, 3, 5, 1, 2])
>>> sorter = jnp.argsort(a)
>>> jnp.searchsorted(a, vals, sorter=sorter)
Array([0, 2, 5, 1, 1], dtype=int32)

The result is equivalent to passing the sorted array:

>>> jnp.searchsorted(jnp.sort(a), vals)
Array([0, 2, 5, 1, 1], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.select(condlist, choicelist, default=0)

Select values based on a series of conditions.

JAX implementation of numpy.select(), implemented in terms of jax.lax.select_n()

Args:
condlist: sequence of array-like conditions. All entries must be mutually

broadcast-compatible.

choicelist: sequence of array-like values to choose. Must have the same length

as condlist, and all entries must be broadcast-compatible with entries of condlist.

default: value to return when every condition is False (default: 0).

Returns:

Array of selected values from choicelist corresponding to the first True entry in condlist at each location.

See also:

  • jax.numpy.where(): select between two values based on a single condition.

  • jax.lax.select_n(): select between N values based on an index.

Examples:
>>> condlist = [
...    jnp.array([False, True, False, False]),
...    jnp.array([True, False, False, False]),
...    jnp.array([False, True, True, False]),
... ]
>>> choicelist = [
...    jnp.array([1, 2, 3, 4]),
...    jnp.array([10, 20, 30, 40]),
...    jnp.array([100, 200, 300, 400]),
... ]
>>> jnp.select(condlist, choicelist, default=0)
Array([ 10,   2, 300,   0], dtype=int32)

This is logically equivalent to the following nested where statement:

>>> default = 0
>>> jnp.where(condlist[0],
...   choicelist[0],
...   jnp.where(condlist[1],
...     choicelist[1],
...     jnp.where(condlist[2],
...       choicelist[2],
...       default)))
Array([ 10,   2, 300,   0], dtype=int32)

However, for efficiency it is implemented in terms of jax.lax.select_n().

Parameters:
Return type:

Array

quchip.declarative.qnp.set_printoptions(*args, **kwargs)[source]

Alias of numpy.set_printoptions().

JAX arrays are printed via NumPy, so NumPy’s printoptions configurations will apply to printed JAX arrays.

See the numpy.set_printoptions() documentation for details on the available options and their meanings.

quchip.declarative.qnp.setdiff1d(ar1, ar2, assume_unique=False, *, size=None, fill_value=None)

Compute the set difference of two 1D arrays.

JAX implementation of numpy.setdiff1d().

Because the size of the output of setdiff1d is data-dependent, the function is not typically compatible with jit() and other JAX transformations. The JAX version adds the optional size argument which must be specified statically for jnp.setdiff1d to be used in such contexts.

Args:

ar1: first array of elements to be differenced. ar2: second array of elements to be differenced. assume_unique: if True, assume the input arrays contain unique values. This allows

a more efficient implementation, but if assume_unique is True and the input arrays contain duplicates, the behavior is undefined. default: False.

size: if specified, return only the first size sorted elements. If there are fewer

elements than size indicates, the return value will be padded with fill_value.

fill_value: when size is specified and there are fewer than the indicated number of

elements, fill the remaining entries fill_value. Defaults to the minimum value.

Returns:

an array containing the set difference of elements in the input array: i.e. the elements in ar1 that are not contained in ar2.

See also:

  • jax.numpy.intersect1d(): the set intersection of two 1D arrays.

  • jax.numpy.setxor1d(): the set XOR of two 1D arrays.

  • jax.numpy.union1d(): the set union of two 1D arrays.

Examples:

Computing the set difference of two arrays:

>>> ar1 = jnp.array([1, 2, 3, 4])
>>> ar2 = jnp.array([3, 4, 5, 6])
>>> jnp.setdiff1d(ar1, ar2)
Array([1, 2], dtype=int32)

Because the output shape is dynamic, this will fail under jit() and other transformations:

>>> jax.jit(jnp.setdiff1d)(ar1, ar2)
Traceback (most recent call last):
   ...
ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: traced array with shape int32[4].
The error occurred while tracing the function setdiff1d at /Users/vanderplas/github/jax-ml/jax/jax/_src/numpy/setops.py:64 for jit. This concrete value was not available in Python because it depends on the value of the argument ar1.

In order to ensure statically-known output shapes, you can pass a static size argument:

>>> jit_setdiff1d = jax.jit(jnp.setdiff1d, static_argnames=['size'])
>>> jit_setdiff1d(ar1, ar2, size=2)
Array([1, 2], dtype=int32)

If size is too small, the difference is truncated:

>>> jit_setdiff1d(ar1, ar2, size=1)
Array([1], dtype=int32)

If size is too large, then the output is padded with fill_value:

>>> jit_setdiff1d(ar1, ar2, size=4, fill_value=0)
Array([1, 2, 0, 0], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.setxor1d(ar1, ar2, assume_unique=False, *, size=None, fill_value=None)

Compute the set-wise xor of elements in two arrays.

JAX implementation of numpy.setxor1d().

Because the size of the output of setxor1d is data-dependent, the function is not compatible with JIT or other JAX transformations.

Args:

ar1: first array of values to intersect. ar2: second array of values to intersect. assume_unique: if True, assume the input arrays contain unique values. This allows

a more efficient implementation, but if assume_unique is True and the input arrays contain duplicates, the behavior is undefined. default: False.

size: if specified, return only the first size sorted elements. If there are fewer

elements than size indicates, the return value will be padded with fill_value, and returned indices will be padded with an out-of-bound index.

fill_value: when size is specified and there are fewer than the indicated number of

elements, fill the remaining entries fill_value. Defaults to the smallest value in the xor result.

Returns:

An array of values that are found in exactly one of the input arrays.

See also:

  • jax.numpy.intersect1d(): the set intersection of two 1D arrays.

  • jax.numpy.union1d(): the set union of two 1D arrays.

  • jax.numpy.setdiff1d(): the set difference of two 1D arrays.

Examples:
>>> ar1 = jnp.array([1, 2, 3, 4])
>>> ar2 = jnp.array([3, 4, 5, 6])
>>> jnp.setxor1d(ar1, ar2)
Array([1, 2, 5, 6], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.shape(a)

Return the shape an array.

JAX implementation of numpy.shape(). Unlike np.shape, this function raises a TypeError if the input is a collection such as a list or tuple.

Args:

a: array-like object, or any object with a shape attribute.

Returns:

An tuple of integers representing the shape of a.

Examples:

Shape for arrays:

>>> x = jnp.arange(10)
>>> jnp.shape(x)
(10,)
>>> y = jnp.ones((2, 3))
>>> jnp.shape(y)
(2, 3)

This also works for scalars:

>>> jnp.shape(3.14)
()

For arrays, this can also be accessed via the jax.Array.shape property:

>>> x.shape
(10,)
Parameters:

a (Array | ndarray | bool | number | bool | int | float | complex | SupportsShape)

Return type:

tuple[int, …]

quchip.declarative.qnp.sign(x, /)

Return an element-wise indication of sign of the input.

JAX implementation of numpy.sign.

The sign of x for real-valued input is:

\[\begin{split}\mathrm{sign}(x) = \begin{cases} 1, & x > 0\\ 0, & x = 0\\ -1, & x < 0 \end{cases}\end{split}\]

For complex valued input, jnp.sign returns a unit vector representing the phase. For generalized case, the sign of x is given by:

\[\begin{split}\mathrm{sign}(x) = \begin{cases} \frac{x}{abs(x)}, & x \ne 0\\ 0, & x = 0 \end{cases}\end{split}\]
Args:

x: input array or scalar.

Returns:

An array with same shape and dtype as x containing the sign indication.

See also:

  • jax.numpy.positive(): Returns element-wise positive values of the input.

  • jax.numpy.negative(): Returns element-wise negative values of the input.

Examples:

For Real-valued inputs:

>>> x = jnp.array([0., -3., 7.])
>>> jnp.sign(x)
Array([ 0., -1.,  1.], dtype=float32)

For complex-inputs:

>>> x1 = jnp.array([1, 3+4j, 5j])
>>> jnp.sign(x1)
Array([1. +0.j , 0.6+0.8j, 0. +1.j ], dtype=complex64)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.signbit(x, /)

Return the sign bit of array elements.

JAX implementation of numpy.signbit.

Args:

x: input array. Complex values are not supported.

Returns:

A boolean array of the same shape as x, containing True where the sign of x is negative, and False otherwise.

See also:

  • jax.numpy.sign(): return the mathematical sign of array elements, i.e. -1, 0, or +1.

Examples:

signbit() on boolean values is always False:

>>> x = jnp.array([True, False])
>>> jnp.signbit(x)
Array([False, False], dtype=bool)

signbit() on integer values is equivalent to x < 0:

>>> x = jnp.array([-2, -1, 0, 1, 2])
>>> jnp.signbit(x)
Array([ True,  True, False, False, False], dtype=bool)

signbit() on floating point values returns the value of the actual sign bit from the float representation, including signed zero:

>>> x = jnp.array([-1.5, -0.0, 0.0, 1.5])
>>> jnp.signbit(x)
Array([ True, True, False, False], dtype=bool)

This also returns the sign bit for special values such as signed NaN and signed infinity:

>>> x = jnp.array([jnp.nan, -jnp.nan, jnp.inf, -jnp.inf])
>>> jnp.signbit(x)
Array([False,  True, False,  True], dtype=bool)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

class quchip.declarative.qnp.signedinteger

Bases: integer

Abstract base class of all signed integer scalar types.

quchip.declarative.qnp.sin(x, /)

Compute a trigonometric sine of each element of input.

JAX implementation of numpy.sin.

Args:

x: array or scalar. Angle in radians.

Returns:

An array containing the sine of each element in x, promotes to inexact dtype.

See also:

  • jax.numpy.cos(): Computes a trigonometric cosine of each element of input.

  • jax.numpy.tan(): Computes a trigonometric tangent of each element of input.

  • jax.numpy.arcsin() and jax.numpy.asin(): Computes the inverse of trigonometric sine of each element of input.

Examples:
>>> pi = jnp.pi
>>> x = jnp.array([pi/4, pi/2, 3*pi/4, pi])
>>> with jnp.printoptions(precision=3, suppress=True):
...   print(jnp.sin(x))
[ 0.707  1.     0.707 -0.   ]
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.sinc(x, /)

Calculate the normalized sinc function.

JAX implementation of numpy.sinc().

The normalized sinc function is given by

\[\mathrm{sinc}(x) = \frac{\sin({\pi x})}{\pi x}\]

where sinc(0) returns the limit value of 1. The sinc function is smooth and infinitely differentiable.

Args:

x : input array; will be promoted to an inexact type.

Returns:

An array of the same shape as x containing the result.

Examples:
>>> x = jnp.array([-1, -0.5, 0, 0.5, 1])
>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.sinc(x)
Array([-0.   ,  0.637,  1.   ,  0.637, -0.   ], dtype=float32)

Compare this to the naive approach to computing the function, which is undefined at zero:

>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.sin(jnp.pi * x) / (jnp.pi * x)
Array([-0.   ,  0.637,    nan,  0.637, -0.   ], dtype=float32)

JAX defines a custom gradient rule for sinc to allow accurate evaluation of the gradient at zero even for higher-order derivatives:

>>> f = jnp.sinc
>>> for i in range(1, 6):
...   f = jax.grad(f)
...   print(f"(d/dx)^{i} f(0.0) = {f(0.0):.2f}")
...
(d/dx)^1 f(0.0) = 0.00
(d/dx)^2 f(0.0) = -3.29
(d/dx)^3 f(0.0) = 0.00
(d/dx)^4 f(0.0) = 19.48
(d/dx)^5 f(0.0) = 0.00
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.single

alias of float32

quchip.declarative.qnp.sinh(x, /)

Calculate element-wise hyperbolic sine of input.

JAX implementation of numpy.sinh.

The hyperbolic sine is defined by:

\[sinh(x) = \frac{e^x - e^{-x}}{2}\]
Args:

x: input array or scalar.

Returns:

An array containing the hyperbolic sine of each element of x, promoting to inexact dtype.

Note:

jnp.sinh is equivalent to computing -1j * jnp.sin(1j * x).

See also:

  • jax.numpy.cosh(): Computes the element-wise hyperbolic cosine of the input.

  • jax.numpy.tanh(): Computes the element-wise hyperbolic tangent of the input.

  • jax.numpy.arcsinh(): Computes the element-wise inverse of hyperbolic sine of the input.

Examples:
>>> x = jnp.array([[-2, 3, 5],
...                [0, -1, 4]])
>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.sinh(x)
Array([[-3.627, 10.018, 74.203],
       [ 0.   , -1.175, 27.29 ]], dtype=float32)
>>> with jnp.printoptions(precision=3, suppress=True):
...   -1j * jnp.sin(1j * x)
Array([[-3.627+0.j, 10.018-0.j, 74.203-0.j],
       [ 0.   -0.j, -1.175+0.j, 27.29 -0.j]],      dtype=complex64, weak_type=True)

For complex-valued input:

>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.sinh(3-2j)
Array(-4.169-9.154j, dtype=complex64, weak_type=True)
>>> with jnp.printoptions(precision=3, suppress=True):
...   -1j * jnp.sin(1j * (3-2j))
Array(-4.169-9.154j, dtype=complex64, weak_type=True)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.size(a, axis=None)

Return number of elements along a given axis.

JAX implementation of numpy.size(). Unlike np.size, this function raises a TypeError if the input is a collection such as a list or tuple.

Args:
a: array-like object, or any object with a size attribute when axis is not

specified, or with a shape attribute when axis is specified.

axis: optional integer along which to count elements. By default, return

the total number of elements.

Returns:

An integer specifying the number of elements in a.

Examples:

Size for arrays:

>>> x = jnp.arange(10)
>>> jnp.size(x)
10
>>> y = jnp.ones((2, 3))
>>> jnp.size(y)
6
>>> jnp.size(y, axis=1)
3

This also works for scalars:

>>> jnp.size(3.14)
1

For arrays, this can also be accessed via the jax.Array.size property:

>>> y.size
6
Parameters:
Return type:

int

quchip.declarative.qnp.sort(a, axis=-1, *, kind=None, order=None, stable=True, descending=False)

Return a sorted copy of an array.

JAX implementation of numpy.sort().

Args:

a: array to sort axis: integer axis along which to sort. Defaults to -1, i.e. the last

axis. If None, then a is flattened before being sorted.

stable: boolean specifying whether a stable sort should be used. Default=True. descending: boolean specifying whether to sort in descending order. Default=False. kind: deprecated; instead specify sort algorithm using stable=True or stable=False. order: not supported by JAX

Returns:

Sorted array of shape a.shape (if axis is an integer) or of shape (a.size,) (if axis is None).

Examples:

Simple 1-dimensional sort

>>> x = jnp.array([1, 3, 5, 4, 2, 1])
>>> jnp.sort(x)
Array([1, 1, 2, 3, 4, 5], dtype=int32)

Sort along the last axis of an array:

>>> x = jnp.array([[2, 1, 3],
...                [4, 3, 6]])
>>> jnp.sort(x, axis=1)
Array([[1, 2, 3],
       [3, 4, 6]], dtype=int32)
See also:

  • jax.numpy.argsort(): return indices of sorted values.

  • jax.numpy.lexsort(): lexicographical sort of multiple arrays.

  • jax.lax.sort(): lower-level function wrapping XLA’s Sort operator.

Parameters:
Return type:

Array

quchip.declarative.qnp.sort_complex(a)

Return a sorted copy of complex array.

JAX implementation of numpy.sort_complex().

Complex numbers are sorted lexicographically, meaning by their real part first, and then by their imaginary part if real parts are equal.

Args:

a: input array. If dtype is not complex, the array will be upcast to complex.

Returns:

A sorted array of the same shape and complex dtype as the input. If a is multi-dimensional, it is sorted along the last axis.

See also:

  • jax.numpy.sort(): Return a sorted copy of an array.

Examples:
>>> a = jnp.array([1+2j, 2+4j, 3-1j, 2+3j])
>>> jnp.sort_complex(a)
Array([1.+2.j, 2.+3.j, 2.+4.j, 3.-1.j], dtype=complex64)

Multi-dimensional arrays are sorted along the last axis:

>>> a = jnp.array([[5, 3, 4],
...                [6, 9, 2]])
>>> jnp.sort_complex(a)
Array([[3.+0.j, 4.+0.j, 5.+0.j],
       [2.+0.j, 6.+0.j, 9.+0.j]], dtype=complex64)
Parameters:

a (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.spacing(x, /)

Return the spacing between x and the next adjacent number.

JAX implementation of numpy.spacing().

Args:

x: real-valued array. Integer or boolean types will be cast to float.

Returns:

Array of same shape as x containing spacing between each entry of x and its closest adjacent value.

See also:

  • jax.numpy.nextafter(): find the next representable value.

Examples:
>>> x = jnp.array([0.0, 0.25, 0.5, 0.75, 1.0], dtype='float32')
>>> jnp.spacing(x)
Array([1.4012985e-45, 2.9802322e-08, 5.9604645e-08, 5.9604645e-08,
      1.1920929e-07], dtype=float32)

For x = 1, the spacing is equal to the eps value given by jax.numpy.finfo:

>>> x = jnp.float32(1)
>>> jnp.spacing(x) == jnp.finfo(x.dtype).eps
Array(True, dtype=bool)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.split(ary, indices_or_sections, axis=0)

Split an array into sub-arrays.

JAX implementation of numpy.split().

Args:

ary: N-dimensional array-like object to split indices_or_sections: either a single integer or a sequence of indices.

  • if indices_or_sections is an integer N, then N must evenly divide ary.shape[axis] and ary will be divided into N equally-sized chunks along axis.

  • if indices_or_sections is a sequence of integers, then these integers specify the boundary between unevenly-sized chunks along axis; see examples below.

axis: the axis along which to split; defaults to 0.

Returns:

A list of arrays. If indices_or_sections is an integer N, then the list is of length N. If indices_or_sections is a sequence seq, then the list is is of length len(seq) + 1.

Examples:

Splitting a 1-dimensional array:

>>> x = jnp.array([1, 2, 3, 4, 5, 6, 7, 8, 9])

Split into three equal sections:

>>> chunks = jnp.split(x, 3)
>>> print(*chunks)
[1 2 3] [4 5 6] [7 8 9]

Split into sections by index:

>>> chunks = jnp.split(x, [2, 7])  # [x[0:2], x[2:7], x[7:]]
>>> print(*chunks)
[1 2] [3 4 5 6 7] [8 9]

Splitting a two-dimensional array along axis 1:

>>> x = jnp.array([[1, 2, 3, 4],
...                [5, 6, 7, 8]])
>>> x1, x2 = jnp.split(x, 2, axis=1)
>>> print(x1)
[[1 2]
 [5 6]]
>>> print(x2)
[[3 4]
 [7 8]]
See also:

  • jax.numpy.array_split(): like split, but allows indices_or_sections to be an integer that does not evenly divide the size of the array.

  • jax.numpy.vsplit(): split vertically, i.e. along axis=0

  • jax.numpy.hsplit(): split horizontally, i.e. along axis=1

  • jax.numpy.dsplit(): split depth-wise, i.e. along axis=2

Parameters:
Return type:

list[Array]

quchip.declarative.qnp.sqrt(x, /)

Calculates element-wise non-negative square root of the input array.

JAX implementation of numpy.sqrt.

Args:

x: input array or scalar.

Returns:

An array containing the non-negative square root of the elements of x.

Note:

  • For real-valued negative inputs, jnp.sqrt produces a nan output.

  • For complex-valued negative inputs, jnp.sqrt produces a complex output.

See also:

  • jax.numpy.square(): Calculates the element-wise square of the input.

  • jax.numpy.power(): Calculates the element-wise base x1 exponential of x2.

Examples:
>>> x = jnp.array([-8-6j, 1j, 4])
>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.sqrt(x)
Array([1.   -3.j   , 0.707+0.707j, 2.   +0.j   ], dtype=complex64)
>>> jnp.sqrt(-1)
Array(nan, dtype=float32, weak_type=True)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.square(x, /)

Calculate element-wise square of the input array.

JAX implementation of numpy.square.

Args:

x: input array or scalar.

Returns:

An array containing the square of the elements of x.

Note:

jnp.square is equivalent to computing jnp.power(x, 2).

See also:

  • jax.numpy.sqrt(): Calculates the element-wise non-negative square root of the input array.

  • jax.numpy.power(): Calculates the element-wise base x1 exponential of x2.

  • jax.lax.integer_pow(): Computes element-wise power \(x^y\), where \(y\) is a fixed integer.

  • jax.numpy.float_power(): Computes the first array raised to the power of second array, element-wise, by promoting to the inexact dtype.

Examples:
>>> x = jnp.array([3, -2, 5.3, 1])
>>> jnp.square(x)
Array([ 9.      ,  4.      , 28.090002,  1.      ], dtype=float32)
>>> jnp.power(x, 2)
Array([ 9.      ,  4.      , 28.090002,  1.      ], dtype=float32)

For integer inputs:

>>> x1 = jnp.array([2, 4, 5, 6])
>>> jnp.square(x1)
Array([ 4, 16, 25, 36], dtype=int32)

For complex-valued inputs:

>>> x2 = jnp.array([1-3j, -1j, 2])
>>> jnp.square(x2)
Array([-8.-6.j, -1.+0.j,  4.+0.j], dtype=complex64)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.squeeze(a, axis=None)

Remove one or more length-1 axes from array

JAX implementation of numpy.sqeeze(), implemented via jax.lax.squeeze().

Args:

a: input array axis: integer or sequence of integers specifying axes to remove. If any specified

axis does not have a length of 1, an error is raised. If not specified, squeeze all length-1 axes in a.

Returns:

copy of a with length-1 axes removed.

Notes:

Unlike numpy.squeeze(), jax.numpy.squeeze() will return a copy rather than a view of the input array. However, under JIT, the compiler will optimize-away such copies when possible, so this doesn’t have performance impacts in practice.

See Also:

  • jax.numpy.expand_dims(): the inverse of squeeze: add dimensions of length 1.

  • jax.Array.squeeze(): equivalent functionality via an array method.

  • jax.lax.squeeze(): equivalent XLA API.

  • jax.numpy.ravel(): flatten an array into a 1D shape.

  • jax.numpy.reshape(): general array reshape.

Examples:
>>> x = jnp.array([[[0]], [[1]], [[2]]])
>>> x.shape
(3, 1, 1)

Squeeze all length-1 dimensions:

>>> jnp.squeeze(x)
Array([0, 1, 2], dtype=int32)
>>> _.shape
(3,)

Equivalent while specifying the axes explicitly:

>>> jnp.squeeze(x, axis=(1, 2))
Array([0, 1, 2], dtype=int32)

Attempting to squeeze a non-unit axis results in an error:

>>> jnp.squeeze(x, axis=0)
Traceback (most recent call last):
  ...
ValueError: cannot select an axis to squeeze out which has size not equal to one, got shape=(3, 1, 1) and dimensions=(0,)

For convenience, this functionality is also available via the jax.Array.squeeze() method:

>>> x.squeeze()
Array([0, 1, 2], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.stack(arrays, axis=0, out=None, dtype=None)

Join arrays along a new axis.

JAX implementation of numpy.stack().

Args:
arrays: a sequence of arrays to stack; each must have the same shape. If a

single array is given it will be treated equivalently to arrays = unstack(arrays), but the implementation will avoid explicit unstacking.

axis: specify the axis along which to stack. out: unused by JAX dtype: optional dtype of the resulting array. If not specified, the dtype

will be determined via type promotion rules described in type-promotion.

Returns:

the stacked result.

See also:

  • jax.numpy.unstack(): inverse of stack.

  • jax.numpy.concatenate(): concatenation along existing axes.

  • jax.numpy.vstack(): stack vertically, i.e. along axis 0.

  • jax.numpy.hstack(): stack horizontally, i.e. along axis 1.

  • jax.numpy.dstack(): stack depth-wise, i.e. along axis 2.

  • jax.numpy.column_stack(): stack columns.

Examples:
>>> x = jnp.array([1, 2, 3])
>>> y = jnp.array([4, 5, 6])
>>> jnp.stack([x, y])
Array([[1, 2, 3],
       [4, 5, 6]], dtype=int32)
>>> jnp.stack([x, y], axis=1)
Array([[1, 4],
       [2, 5],
       [3, 6]], dtype=int32)

unstack() performs the inverse operation:

>>> arr = jnp.stack([x, y], axis=1)
>>> x, y = jnp.unstack(arr, axis=1)
>>> x
Array([1, 2, 3], dtype=int32)
>>> y
Array([4, 5, 6], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.std(a, axis=None, dtype=None, out=None, ddof=0, keepdims=False, *, where=None, correction=None)

Compute the standard deviation along a given axis.

JAX implementation of numpy.std().

Args:

a: input array. axis: optional, int or sequence of ints, default=None. Axis along which the

standard deviation is computed. If None, standard deviaiton is computed along all the axes.

dtype: The type of the output array. Default=None. ddof: int, default=0. Degrees of freedom. The divisor in the standard deviation

computation is N-ddof, N is number of elements along given axis.

keepdims: bool, default=False. If true, reduced axes are left in the result

with size 1.

where: optional, boolean array, default=None. The elements to be used in the

standard deviation. Array should be broadcast compatible to the input.

correction: int or float, default=None. Alternative name for ddof.

Both ddof and correction can’t be provided simultaneously.

out: Unused by JAX.

Returns:

An array of the standard deviation along the given axis.

See also:

  • jax.numpy.var(): Compute the variance of array elements over given axis.

  • jax.numpy.mean(): Compute the mean of array elements over a given axis.

  • jax.numpy.nanvar(): Compute the variance along a given axis, ignoring NaNs values.

  • jax.numpy.nanstd(): Computed the standard deviation of a given axis, ignoring NaN values.

Examples:

By default, jnp.std computes the standard deviation along all axes.

>>> x = jnp.array([[1, 3, 4, 2],
...                [4, 2, 5, 3],
...                [5, 4, 2, 3]])
>>> with jnp.printoptions(precision=2, suppress=True):
...   jnp.std(x)
Array(1.21, dtype=float32)

If axis=0, computes along axis 0.

>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jnp.std(x, axis=0))
[1.7  0.82 1.25 0.47]

To preserve the dimensions of input, you can set keepdims=True.

>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jnp.std(x, axis=0, keepdims=True))
[[1.7  0.82 1.25 0.47]]

If ddof=1:

>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jnp.std(x, axis=0, keepdims=True, ddof=1))
[[2.08 1.   1.53 0.58]]

To include specific elements of the array to compute standard deviation, you can use where.

>>> where = jnp.array([[1, 0, 1, 0],
...                    [0, 1, 0, 1],
...                    [1, 1, 1, 0]], dtype=bool)
>>> jnp.std(x, axis=0, keepdims=True, where=where)
Array([[2., 1., 1., 0.]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.sum(a, axis=None, dtype=None, out=None, keepdims=False, initial=None, where=None, promote_integers=True)

Sum of the elements of the array over a given axis.

JAX implementation of numpy.sum().

Args:

a: Input array. axis: int or array, default=None. Axis along which the sum to be computed.

If None, the sum is computed along all the axes.

dtype: The type of the output array. Default=None. out: Unused by JAX keepdims: bool, default=False. If true, reduced axes are left in the result

with size 1.

initial: int or array, Default=None. Initial value for the sum. where: int or array, default=None. The elements to be used in the sum. Array

should be broadcast compatible to the input.

promote_integersbool, default=True. If True, then integer inputs will be

promoted to the widest available integer dtype, following numpy’s behavior. If False, the result will have the same dtype as the input. promote_integers is ignored if dtype is specified.

Returns:

An array of the sum along the given axis.

See also:

  • jax.numpy.prod(): Compute the product of array elements over a given axis.

  • jax.numpy.max(): Compute the maximum of array elements over given axis.

  • jax.numpy.min(): Compute the minimum of array elements over given axis.

Examples:

By default, the sum is computed along all the axes.

>>> x = jnp.array([[1, 3, 4, 2],
...                [5, 2, 6, 3],
...                [8, 1, 3, 9]])
>>> jnp.sum(x)
Array(47, dtype=int32)

If axis=1, the sum is computed along axis 1.

>>> jnp.sum(x, axis=1)
Array([10, 16, 21], dtype=int32)

If keepdims=True, ndim of the output is equal to that of the input.

>>> jnp.sum(x, axis=1, keepdims=True)
Array([[10],
       [16],
       [21]], dtype=int32)

To include only specific elements in the sum, you can use where.

>>> where=jnp.array([[0, 0, 1, 0],
...                  [0, 0, 1, 1],
...                  [1, 1, 1, 0]], dtype=bool)
>>> jnp.sum(x, axis=1, keepdims=True, where=where)
Array([[ 4],
       [ 9],
       [12]], dtype=int32)
>>> where=jnp.array([[False],
...                  [False],
...                  [False]])
>>> jnp.sum(x, axis=0, keepdims=True, where=where)
Array([[0, 0, 0, 0]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.swapaxes(a, axis1, axis2)

Swap two axes of an array.

JAX implementation of numpy.swapaxes(), implemented in terms of jax.lax.transpose().

Args:

a: input array axis1: index of first axis axis2: index of second axis

Returns:

Copy of a with specified axes swapped.

Notes:

Unlike numpy.swapaxes(), jax.numpy.swapaxes() will return a copy rather than a view of the input array. However, under JIT, the compiler will optimize away such copies when possible, so this doesn’t have performance impacts in practice.

See Also:

  • jax.numpy.moveaxis(): move a single axis of an array.

  • jax.numpy.rollaxis(): older API for moveaxis.

  • jax.lax.transpose(): more general axes permutations.

  • jax.Array.swapaxes(): same functionality via an array method.

Examples:
>>> a = jnp.ones((2, 3, 4, 5))
>>> jnp.swapaxes(a, 1, 3).shape
(2, 5, 4, 3)

Equivalent output via the swapaxes array method:

>>> a.swapaxes(1, 3).shape
(2, 5, 4, 3)

Equivalent output via transpose():

>>> a.transpose(0, 3, 2, 1).shape
(2, 5, 4, 3)
Parameters:
Return type:

Array

quchip.declarative.qnp.take(a, indices, axis=None, out=None, mode=None, unique_indices=False, indices_are_sorted=False, fill_value=None)

Take elements from an array.

JAX implementation of numpy.take(), implemented in terms of jax.lax.gather(). JAX’s behavior differs from NumPy in the case of out-of-bound indices; see the mode parameter below.

Args:

a: array from which to take values. indices: N-dimensional array of integer indices of values to take from the array. axis: the axis along which to take values. If not specified, the array will

be flattened before indexing is applied.

mode: Out-of-bounds indexing mode, either "fill" or "clip". The default

mode="fill" returns invalid values (e.g. NaN) for out-of bounds indices; the fill_value argument gives control over this value. For more discussion of mode options, see jax.numpy.ndarray.at.

fill_value: The fill value to return for out-of-bounds slices when mode is ‘fill’.

Ignored otherwise. Defaults to NaN for inexact types, the largest negative value for signed types, the largest positive value for unsigned types, and True for booleans.

unique_indices: If True, the implementation will assume that the indices are unique,

which can result in more efficient execution on some backends. If set to True and indices are not unique, the output is undefined.

indices_are_sortedIf True, the implementation will assume that the indices are

sorted in ascending order, which can lead to more efficient execution on some backends. If set to True and indices are not sorted, the output is undefined.

Returns:

Array of values extracted from a.

See also:

  • jax.numpy.ndarray.at: take values via indexing syntax.

  • jax.numpy.take_along_axis(): take values along an axis

Examples:
>>> x = jnp.array([[1., 2., 3.],
...                [4., 5., 6.]])
>>> indices = jnp.array([2, 0])

Passing no axis results in indexing into the flattened array:

>>> jnp.take(x, indices)
Array([3., 1.], dtype=float32)
>>> x.ravel()[indices]  # equivalent indexing syntax
Array([3., 1.], dtype=float32)

Passing an axis results ind applying the index to every subarray along the axis:

>>> jnp.take(x, indices, axis=1)
Array([[3., 1.],
       [6., 4.]], dtype=float32)
>>> x[:, indices]  # equivalent indexing syntax
Array([[3., 1.],
       [6., 4.]], dtype=float32)

Out-of-bound indices fill with invalid values. For float inputs, this is NaN:

>>> jnp.take(x, indices, axis=0)
Array([[nan, nan, nan],
       [ 1.,  2.,  3.]], dtype=float32)
>>> x.at[indices].get(mode='fill', fill_value=jnp.nan)  # equivalent indexing syntax
Array([[nan, nan, nan],
       [ 1.,  2.,  3.]], dtype=float32)

This default out-of-bound behavior can be adjusted using the mode parameter, for example, we can instead clip to the last valid value:

>>> jnp.take(x, indices, axis=0, mode='clip')
Array([[4., 5., 6.],
       [1., 2., 3.]], dtype=float32)
>>> x.at[indices].get(mode='clip')  # equivalent indexing syntax
Array([[4., 5., 6.],
       [1., 2., 3.]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.take_along_axis(arr, indices, axis, mode=None, fill_value=None)

Take elements from an array.

JAX implementation of numpy.take_along_axis(), implemented in terms of jax.lax.gather(). JAX’s behavior differs from NumPy in the case of out-of-bound indices; see the mode parameter below.

Args:

a: array from which to take values. indices: array of integer indices. If axis is None, must be one-dimensional.

If axis is not None, must have a.ndim == indices.ndim, and a must be broadcast-compatible with indices along dimensions other than axis.

axis: the axis along which to take values. If not specified, the array will

be flattened before indexing is applied.

mode: Out-of-bounds indexing mode, either "fill" or "clip". The default

mode="fill" returns invalid values (e.g. NaN) for out-of bounds indices. For more discussion of mode options, see jax.numpy.ndarray.at.

Returns:

Array of values extracted from a.

See also:

  • jax.numpy.ndarray.at: take values via indexing syntax.

  • jax.numpy.take(): take the same indices along every axis slice.

Examples:
>>> x = jnp.array([[1., 2., 3.],
...                [4., 5., 6.]])
>>> indices = jnp.array([[0, 2],
...                      [1, 0]])
>>> jnp.take_along_axis(x, indices, axis=1)
Array([[1., 3.],
       [5., 4.]], dtype=float32)
>>> x[jnp.arange(2)[:, None], indices]  # equivalent via indexing syntax
Array([[1., 3.],
       [5., 4.]], dtype=float32)

Out-of-bound indices fill with invalid values. For float inputs, this is NaN:

>>> indices = jnp.array([[1, 0, 2]])
>>> jnp.take_along_axis(x, indices, axis=0)
Array([[ 4.,  2., nan]], dtype=float32)
>>> x.at[indices, jnp.arange(3)].get(
...     mode='fill', fill_value=jnp.nan)  # equivalent via indexing syntax
Array([[ 4.,  2., nan]], dtype=float32)

take_along_axis is helpful for extracting values from multi-dimensional argsorts and arg reductions. For, here we compute argsort() indices along an axis, and use take_along_axis to construct the sorted array:

>>> x = jnp.array([[5, 3, 4],
...                [2, 7, 6]])
>>> indices = jnp.argsort(x, axis=1)
>>> indices
Array([[1, 2, 0],
       [0, 2, 1]], dtype=int32)
>>> jnp.take_along_axis(x, indices, axis=1)
Array([[3, 4, 5],
       [2, 6, 7]], dtype=int32)

Similarly, we can use argmin() with keepdims=True and use take_along_axis to extract the minimum value:

>>> idx = jnp.argmin(x, axis=1, keepdims=True)
>>> idx
Array([[1],
       [0]], dtype=int32)
>>> jnp.take_along_axis(x, idx, axis=1)
Array([[3],
       [2]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.tan(x, /)

Compute a trigonometric tangent of each element of input.

JAX implementation of numpy.tan.

Args:

x: scalar or array. Angle in radians.

Returns:

An array containing the tangent of each element in x, promotes to inexact dtype.

See also:

  • jax.numpy.sin(): Computes a trigonometric sine of each element of input.

  • jax.numpy.cos(): Computes a trigonometric cosine of each element of input.

  • jax.numpy.arctan() and jax.numpy.atan(): Computes the inverse of trigonometric tangent of each element of input.

Examples:
>>> pi = jnp.pi
>>> x = jnp.array([0, pi/6, pi/4, 3*pi/4, 5*pi/6])
>>> with jnp.printoptions(precision=3, suppress=True):
...   print(jnp.tan(x))
[ 0.     0.577  1.    -1.    -0.577]
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.tanh(x, /)

Calculate element-wise hyperbolic tangent of input.

JAX implementation of numpy.tanh.

The hyperbolic tangent is defined by:

\[tanh(x) = \frac{sinh(x)}{cosh(x)} = \frac{e^x - e^{-x}}{e^x + e^{-x}}\]
Args:

x: input array or scalar.

Returns:

An array containing the hyperbolic tangent of each element of x, promoting to inexact dtype.

Note:

jnp.tanh is equivalent to computing -1j * jnp.tan(1j * x).

See also:

  • jax.numpy.sinh(): Computes the element-wise hyperbolic sine of the input.

  • jax.numpy.cosh(): Computes the element-wise hyperbolic cosine of the input.

  • jax.numpy.arctanh(): Computes the element-wise inverse of hyperbolic tangent of the input.

Examples:
>>> x = jnp.array([[-1, 0, 1],
...                [3, -2, 5]])
>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.tanh(x)
Array([[-0.762,  0.   ,  0.762],
       [ 0.995, -0.964,  1.   ]], dtype=float32)
>>> with jnp.printoptions(precision=3, suppress=True):
...   -1j * jnp.tan(1j * x)
Array([[-0.762+0.j,  0.   -0.j,  0.762-0.j],
       [ 0.995-0.j, -0.964+0.j,  1.   -0.j]],      dtype=complex64, weak_type=True)

For complex-valued input:

>>> with jnp.printoptions(precision=3, suppress=True):
...   jnp.tanh(2-5j)
Array(1.031+0.021j, dtype=complex64, weak_type=True)
>>> with jnp.printoptions(precision=3, suppress=True):
...   -1j * jnp.tan(1j * (2-5j))
Array(1.031+0.021j, dtype=complex64, weak_type=True)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

quchip.declarative.qnp.tensordot(a, b, axes=2, *, precision=None, preferred_element_type=None)

Compute the tensor dot product of two N-dimensional arrays.

JAX implementation of numpy.linalg.tensordot().

Args:

a: N-dimensional array b: M-dimensional array axes: integer or tuple of sequences of integers. If an integer k, then

sum over the last k axes of a and the first k axes of b, in order. If a tuple, then axes[0] specifies the axes of a and axes[1] specifies the axes of b.

precision: either None (default), which means the default precision for

the backend, a Precision enum value (Precision.DEFAULT, Precision.HIGH or Precision.HIGHEST) or a tuple of two such values indicating precision of a and b.

preferred_element_type: either None (default), which means the default

accumulation type for the input types, or a datatype, indicating to accumulate results to and return a result with that datatype.

Returns:

array containing the tensor dot product of the inputs

See also:

  • jax.numpy.einsum(): NumPy API for more general tensor contractions.

  • jax.lax.dot_general(): XLA API for more general tensor contractions.

Examples:
>>> x1 = jnp.arange(24.).reshape(2, 3, 4)
>>> x2 = jnp.ones((3, 4, 5))
>>> jnp.tensordot(x1, x2)
Array([[ 66.,  66.,  66.,  66.,  66.],
       [210., 210., 210., 210., 210.]], dtype=float32)

Equivalent result when specifying the axes as explicit sequences:

>>> jnp.tensordot(x1, x2, axes=([1, 2], [0, 1]))
Array([[ 66.,  66.,  66.,  66.,  66.],
       [210., 210., 210., 210., 210.]], dtype=float32)

Equivalent result via einsum():

>>> jnp.einsum('ijk,jkm->im', x1, x2)
Array([[ 66.,  66.,  66.,  66.,  66.],
       [210., 210., 210., 210., 210.]], dtype=float32)

Setting axes=1 for two-dimensional inputs is equivalent to a matrix multiplication:

>>> x1 = jnp.array([[1, 2],
...                 [3, 4]])
>>> x2 = jnp.array([[1, 2, 3],
...                 [4, 5, 6]])
>>> jnp.linalg.tensordot(x1, x2, axes=1)
Array([[ 9, 12, 15],
       [19, 26, 33]], dtype=int32)
>>> x1 @ x2
Array([[ 9, 12, 15],
       [19, 26, 33]], dtype=int32)

Setting axes=0 for one-dimensional inputs is equivalent to outer():

>>> x1 = jnp.array([1, 2])
>>> x2 = jnp.array([1, 2, 3])
>>> jnp.linalg.tensordot(x1, x2, axes=0)
Array([[1, 2, 3],
       [2, 4, 6]], dtype=int32)
>>> jnp.outer(x1, x2)
Array([[1, 2, 3],
       [2, 4, 6]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.tile(A, reps)

Construct an array by repeating A along specified dimensions.

JAX implementation of numpy.tile().

If A is an array of shape (d1, d2, ..., dn) and reps is a sequence of integers, the resulting array will have a shape of (reps[0] * d1, reps[1] * d2, ..., reps[n] * dn), with A tiled along each dimension.

Args:

A: input array to be repeated. Can be of any shape or dimension. reps: specifies the number of repetitions along each axis.

Returns:

a new array where the input array has been repeated according to reps.

See also:

  • jax.numpy.repeat(): Construct an array from repeated elements.

  • jax.numpy.broadcast_to(): Broadcast an array to a specified shape.

Examples:
>>> arr = jnp.array([1, 2])
>>> jnp.tile(arr, 2)
Array([1, 2, 1, 2], dtype=int32)
>>> arr = jnp.array([[1, 2],
...                  [3, 4,]])
>>> jnp.tile(arr, (2, 1))
Array([[1, 2],
       [3, 4],
       [1, 2],
       [3, 4]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.trace(a, offset=0, axis1=0, axis2=1, dtype=None, out=None)

Calculate sum of the diagonal of input along the given axes.

JAX implementation of numpy.trace().

Args:

a: input array. Must have a.ndim >= 2. offset: optional, int, default=0. Diagonal offset from the main diagonal.

Can be positive or negative.

axis1: optional, default=0. The first axis along which to take the sum of

diagonal. Must be a static integer value.

axis2: optional, default=1. The second axis along which to take the sum of

diagonal. Must be a static integer value.

dtype: optional. The dtype of the output array. Should be provided as static

argument in JIT compilation.

out: Not used by JAX.

Returns:

An array of dimension x.ndim-2 containing the sum of the diagonal elements along axes (axis1, axis2)

See also:

  • jax.numpy.diag(): Returns the specified diagonal or constructs a diagonal array

  • jax.numpy.diagonal(): Returns the specified diagonal of an array.

  • jax.numpy.diagflat(): Returns a 2-D array with the flattened input array laid out on the diagonal.

Examples:
>>> x = jnp.arange(1, 9).reshape(2, 2, 2)
>>> x
Array([[[1, 2],
        [3, 4]],

       [[5, 6],
        [7, 8]]], dtype=int32)
>>> jnp.trace(x)
Array([ 8, 10], dtype=int32)
>>> jnp.trace(x, offset=1)
Array([3, 4], dtype=int32)
>>> jnp.trace(x, axis1=1, axis2=2)
Array([ 5, 13], dtype=int32)
>>> jnp.trace(x, offset=1, axis1=1, axis2=2)
Array([2, 6], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.transpose(a, axes=None)

Return a transposed version of an N-dimensional array.

JAX implementation of numpy.transpose(), implemented in terms of jax.lax.transpose().

Args:

a: input array axes: optionally specify the permutation using a length-a.ndim sequence of integers

i satisfying 0 <= i < a.ndim. Defaults to range(a.ndim)[::-1], i.e. reverses the order of all axes.

Returns:

transposed copy of the array.

See Also:

  • jax.Array.transpose(): equivalent function via an Array method.

  • jax.Array.T: equivalent function via an Array property.

  • jax.numpy.matrix_transpose(): transpose the last two axes of an array. This is suitable for working with batched 2D matrices.

  • jax.numpy.swapaxes(): swap any two axes in an array.

  • jax.numpy.moveaxis(): move an axis to another position in the array.

Note:

Unlike numpy.transpose(), jax.numpy.transpose() will return a copy rather than a view of the input array. However, under JIT, the compiler will optimize-away such copies when possible, so this doesn’t have performance impacts in practice.

Examples:

For a 1D array, the transpose is the identity:

>>> x = jnp.array([1, 2, 3, 4])
>>> jnp.transpose(x)
Array([1, 2, 3, 4], dtype=int32)

For a 2D array, the transpose is a matrix transpose:

>>> x = jnp.array([[1, 2],
...                [3, 4]])
>>> jnp.transpose(x)
Array([[1, 3],
       [2, 4]], dtype=int32)

For an N-dimensional array, the transpose reverses the order of the axes:

>>> x = jnp.zeros(shape=(3, 4, 5))
>>> jnp.transpose(x).shape
(5, 4, 3)

The axes argument can be specified to change this default behavior:

>>> jnp.transpose(x, (0, 2, 1)).shape
(3, 5, 4)

Since swapping the last two axes is a common operation, it can be done via its own API, jax.numpy.matrix_transpose():

>>> jnp.matrix_transpose(x).shape
(3, 5, 4)

For convenience, transposes may also be performed using the jax.Array.transpose() method or the jax.Array.T property:

>>> x = jnp.array([[1, 2],
...                [3, 4]])
>>> x.transpose()
Array([[1, 3],
       [2, 4]], dtype=int32)
>>> x.T
Array([[1, 3],
       [2, 4]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.trapezoid(y, x=None, dx=1.0, axis=-1)

Integrate along the given axis using the composite trapezoidal rule.

JAX implementation of numpy.trapezoid()

The trapezoidal rule approximates the integral under a curve by summing the areas of trapezoids formed between adjacent data points.

Args:

y: array of data to integrate. x: optional array of sample points corresponding to the y values. If not

provided, x defaults to equally spaced with spacing given by dx.

dx: The spacing between sample points when x is None (default: 1.0). axis: The axis along which to integrate (default: -1)

Returns:

The definite integral approximated by the trapezoidal rule.

Examples:

Integrate over a regular grid, with spacing 1.0:

>>> y = jnp.array([1, 2, 3, 2, 3, 2, 1])
>>> jnp.trapezoid(y, dx=1.0)
Array(13., dtype=float32)

Integrate over an irregular grid:

>>> x = jnp.array([0, 2, 5, 7, 10, 15, 20])
>>> jnp.trapezoid(y, x)
Array(43., dtype=float32)

Approximate \(\int_0^{2\pi} \sin^2(x)dx\), which equals \(\pi\):

>>> x = jnp.linspace(0, 2 * jnp.pi, 1000)
>>> y = jnp.sin(x) ** 2
>>> result = jnp.trapezoid(y, x)
>>> jnp.allclose(result, jnp.pi)
Array(True, dtype=bool)
Parameters:
Return type:

Array

quchip.declarative.qnp.tri(N, M=None, k=0, dtype=None)

Return an array with ones on and below the diagonal and zeros elsewhere.

JAX implementation of numpy.tri()

Args:

N: int. Dimension of the rows of the returned array. M: optional, int. Dimension of the columns of the returned array. If not

specified, then M = N.

k: optional, int, default=0. Specifies the sub-diagonal on and below which

the array is filled with ones. k=0 refers to main diagonal, k<0 refers to sub-diagonal below the main diagonal and k>0 refers to sub-diagonal above the main diagonal.

dtype: optional, data type of the returned array. The default type is float.

Returns:

An array of shape (N, M) containing the lower triangle with elements below the sub-diagonal specified by k are set to one and zero elsewhere.

See also:

  • jax.numpy.tril(): Returns a lower triangle of an array.

  • jax.numpy.triu(): Returns an upper triangle of an array.

Examples:
>>> jnp.tri(3)
Array([[1., 0., 0.],
       [1., 1., 0.],
       [1., 1., 1.]], dtype=float32)

When M is not equal to N:

>>> jnp.tri(3, 4)
Array([[1., 0., 0., 0.],
       [1., 1., 0., 0.],
       [1., 1., 1., 0.]], dtype=float32)

when k>0:

>>> jnp.tri(3, k=1)
Array([[1., 1., 0.],
       [1., 1., 1.],
       [1., 1., 1.]], dtype=float32)

When k<0:

>>> jnp.tri(3, 4, k=-1)
Array([[0., 0., 0., 0.],
       [1., 0., 0., 0.],
       [1., 1., 0., 0.]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.tril(m, k=0)

Return lower triangle of an array.

JAX implementation of numpy.tril()

Args:

m: input array. Must have m.ndim >= 2. k: k: optional, int, default=0. Specifies the sub-diagonal above which the

elements of the array are set to zero. k=0 refers to main diagonal, k<0 refers to sub-diagonal below the main diagonal and k>0 refers to sub-diagonal above the main diagonal.

Returns:

An array with same shape as input containing the lower triangle of the given array with elements above the sub-diagonal specified by k are set to zero.

See also:

  • jax.numpy.triu(): Returns an upper triangle of an array.

  • jax.numpy.tri(): Returns an array with ones on and below the diagonal and zeros elsewhere.

Examples:
>>> x = jnp.array([[1, 2, 3, 4],
...                [5, 6, 7, 8],
...                [9, 10, 11, 12]])
>>> jnp.tril(x)
Array([[ 1,  0,  0,  0],
       [ 5,  6,  0,  0],
       [ 9, 10, 11,  0]], dtype=int32)
>>> jnp.tril(x, k=1)
Array([[ 1,  2,  0,  0],
       [ 5,  6,  7,  0],
       [ 9, 10, 11, 12]], dtype=int32)
>>> jnp.tril(x, k=-1)
Array([[ 0,  0,  0,  0],
       [ 5,  0,  0,  0],
       [ 9, 10,  0,  0]], dtype=int32)

When m.ndim > 2, jnp.tril operates batch-wise on the trailing axes.

>>> x1 = jnp.array([[[1, 2],
...                  [3, 4]],
...                 [[5, 6],
...                  [7, 8]]])
>>> jnp.tril(x1)
Array([[[1, 0],
        [3, 4]],

       [[5, 0],
        [7, 8]]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.tril_indices(n, k=0, m=None)

Return the indices of lower triangle of an array of size (n, m).

JAX implementation of numpy.tril_indices().

Args:

n: int. Number of rows of the array for which the indices are returned. k: optional, int, default=0. Specifies the sub-diagonal on and below which

the indices of lower triangle are returned. k=0 refers to main diagonal, k<0 refers to sub-diagonal below the main diagonal and k>0 refers to sub-diagonal above the main diagonal.

m: optional, int. Number of columns of the array for which the indices are

returned. If not specified, then m = n.

Returns:

A tuple of two arrays containing the indices of the lower triangle, one along each axis.

See also:

  • jax.numpy.triu_indices(): Returns the indices of upper triangle of an array of size (n, m).

  • jax.numpy.triu_indices_from(): Returns the indices of upper triangle of a given array.

  • jax.numpy.tril_indices_from(): Returns the indices of lower triangle of a given array.

Examples:

If only n is provided in input, the indices of lower triangle of an array of size (n, n) array are returned.

>>> jnp.tril_indices(3)
(Array([0, 1, 1, 2, 2, 2], dtype=int32), Array([0, 0, 1, 0, 1, 2], dtype=int32))

If both n and m are provided in input, the indices of lower triangle of an (n, m) array are returned.

>>> jnp.tril_indices(3, m=2)
(Array([0, 1, 1, 2, 2], dtype=int32), Array([0, 0, 1, 0, 1], dtype=int32))

If k = 1, the indices on and below the first sub-diagonal above the main diagonal are returned.

>>> jnp.tril_indices(3, k=1)
(Array([0, 0, 1, 1, 1, 2, 2, 2], dtype=int32), Array([0, 1, 0, 1, 2, 0, 1, 2], dtype=int32))

If k = -1, the indices on and below the first sub-diagonal below the main diagonal are returned.

>>> jnp.tril_indices(3, k=-1)
(Array([1, 2, 2], dtype=int32), Array([0, 0, 1], dtype=int32))
Parameters:
Return type:

tuple[Array, Array]

quchip.declarative.qnp.tril_indices_from(arr, k=0)

Return the indices of lower triangle of a given array.

JAX implementation of numpy.tril_indices_from().

Args:

arr: input array. Must have arr.ndim == 2. k: optional, int, default=0. Specifies the sub-diagonal on and below which

the indices of upper triangle are returned. k=0 refers to main diagonal, k<0 refers to sub-diagonal below the main diagonal and k>0 refers to sub-diagonal above the main diagonal.

Returns:

A tuple of two arrays containing the indices of the lower triangle, one along each axis.

See also:

  • jax.numpy.triu_indices_from(): Returns the indices of upper triangle of a given array.

  • jax.numpy.tril_indices(): Returns the indices of lower triangle of an array of size (n, m).

  • jax.numpy.tril(): Returns a lower triangle of an array

Examples:
>>> arr = jnp.array([[1, 2, 3],
...                  [4, 5, 6],
...                  [7, 8, 9]])
>>> jnp.tril_indices_from(arr)
(Array([0, 1, 1, 2, 2, 2], dtype=int32), Array([0, 0, 1, 0, 1, 2], dtype=int32))

Elements indexed by jnp.tril_indices_from correspond to those in the output of jnp.tril.

>>> ind = jnp.tril_indices_from(arr)
>>> arr[ind]
Array([1, 4, 5, 7, 8, 9], dtype=int32)
>>> jnp.tril(arr)
Array([[1, 0, 0],
       [4, 5, 0],
       [7, 8, 9]], dtype=int32)

When k > 0:

>>> jnp.tril_indices_from(arr, k=1)
(Array([0, 0, 1, 1, 1, 2, 2, 2], dtype=int32), Array([0, 1, 0, 1, 2, 0, 1, 2], dtype=int32))

When k < 0:

>>> jnp.tril_indices_from(arr, k=-1)
(Array([1, 2, 2], dtype=int32), Array([0, 0, 1], dtype=int32))
Parameters:
Return type:

tuple[Array, Array]

quchip.declarative.qnp.trim_zeros(filt, trim='fb')

Trim leading and/or trailing zeros of the input array.

JAX implementation of numpy.trim_zeros().

Args:

filt: input array. Must have filt.ndim == 1. trim: string, optional, default = fb. Specifies from which end the input

is trimmed.

  • f - trims only the leading zeros.

  • b - trims only the trailing zeros.

  • fb - trims both leading and trailing zeros.

Returns:

An array containing the trimmed input with same dtype as filt.

Examples:
>>> x = jnp.array([0, 0, 2, 0, 1, 4, 3, 0, 0, 0])
>>> jnp.trim_zeros(x)
Array([2, 0, 1, 4, 3], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.triu(m, k=0)

Return upper triangle of an array.

JAX implementation of numpy.triu()

Args:

m: input array. Must have m.ndim >= 2. k: optional, int, default=0. Specifies the sub-diagonal below which the

elements of the array are set to zero. k=0 refers to main diagonal, k<0 refers to sub-diagonal below the main diagonal and k>0 refers to sub-diagonal above the main diagonal.

Returns:

An array with same shape as input containing the upper triangle of the given array with elements below the sub-diagonal specified by k are set to zero.

See also:

  • jax.numpy.tril(): Returns a lower triangle of an array.

  • jax.numpy.tri(): Returns an array with ones on and below the diagonal and zeros elsewhere.

Examples:
>>> x = jnp.array([[1, 2, 3],
...                [4, 5, 6],
...                [7, 8, 9],
...                [10, 11, 12]])
>>> jnp.triu(x)
Array([[1, 2, 3],
       [0, 5, 6],
       [0, 0, 9],
       [0, 0, 0]], dtype=int32)
>>> jnp.triu(x, k=1)
Array([[0, 2, 3],
       [0, 0, 6],
       [0, 0, 0],
       [0, 0, 0]], dtype=int32)
>>> jnp.triu(x, k=-1)
Array([[ 1,  2,  3],
       [ 4,  5,  6],
       [ 0,  8,  9],
       [ 0,  0, 12]], dtype=int32)

When m.ndim > 2, jnp.triu operates batch-wise on the trailing axes.

>>> x1 = jnp.array([[[1, 2],
...                  [3, 4]],
...                 [[5, 6],
...                  [7, 8]]])
>>> jnp.triu(x1)
Array([[[1, 2],
        [0, 4]],

       [[5, 6],
        [0, 8]]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.triu_indices(n, k=0, m=None)

Return the indices of upper triangle of an array of size (n, m).

JAX implementation of numpy.triu_indices().

Args:

n: int. Number of rows of the array for which the indices are returned. k: optional, int, default=0. Specifies the sub-diagonal on and above which

the indices of upper triangle are returned. k=0 refers to main diagonal, k<0 refers to sub-diagonal below the main diagonal and k>0 refers to sub-diagonal above the main diagonal.

m: optional, int. Number of columns of the array for which the indices are

returned. If not specified, then m = n.

Returns:

A tuple of two arrays containing the indices of the upper triangle, one along each axis.

See also:

  • jax.numpy.tril_indices(): Returns the indices of lower triangle of an array of size (n, m).

  • jax.numpy.triu_indices_from(): Returns the indices of upper triangle of a given array.

  • jax.numpy.tril_indices_from(): Returns the indices of lower triangle of a given array.

Examples:

If only n is provided in input, the indices of upper triangle of an array of size (n, n) array are returned.

>>> jnp.triu_indices(3)
(Array([0, 0, 0, 1, 1, 2], dtype=int32), Array([0, 1, 2, 1, 2, 2], dtype=int32))

If both n and m are provided in input, the indices of upper triangle of an (n, m) array are returned.

>>> jnp.triu_indices(3, m=2)
(Array([0, 0, 1], dtype=int32), Array([0, 1, 1], dtype=int32))

If k = 1, the indices on and above the first sub-diagonal above the main diagonal are returned.

>>> jnp.triu_indices(3, k=1)
(Array([0, 0, 1], dtype=int32), Array([1, 2, 2], dtype=int32))

If k = -1, the indices on and above the first sub-diagonal below the main diagonal are returned.

>>> jnp.triu_indices(3, k=-1)
(Array([0, 0, 0, 1, 1, 1, 2, 2], dtype=int32), Array([0, 1, 2, 0, 1, 2, 1, 2], dtype=int32))
Parameters:
Return type:

tuple[Array, Array]

quchip.declarative.qnp.triu_indices_from(arr, k=0)

Return the indices of upper triangle of a given array.

JAX implementation of numpy.triu_indices_from().

Args:

arr: input array. Must have arr.ndim == 2. k: optional, int, default=0. Specifies the sub-diagonal on and above which

the indices of upper triangle are returned. k=0 refers to main diagonal, k<0 refers to sub-diagonal below the main diagonal and k>0 refers to sub-diagonal above the main diagonal.

Returns:

A tuple of two arrays containing the indices of the upper triangle, one along each axis.

See also:

  • jax.numpy.tril_indices_from(): Returns the indices of lower triangle of a given array.

  • jax.numpy.triu_indices(): Returns the indices of upper triangle of an array of size (n, m).

  • jax.numpy.triu(): Return an upper triangle of an array.

Examples:
>>> arr = jnp.array([[1, 2, 3],
...                  [4, 5, 6],
...                  [7, 8, 9]])
>>> jnp.triu_indices_from(arr)
(Array([0, 0, 0, 1, 1, 2], dtype=int32), Array([0, 1, 2, 1, 2, 2], dtype=int32))

Elements indexed by jnp.triu_indices_from correspond to those in the output of jnp.triu.

>>> ind = jnp.triu_indices_from(arr)
>>> arr[ind]
Array([1, 2, 3, 5, 6, 9], dtype=int32)
>>> jnp.triu(arr)
Array([[1, 2, 3],
       [0, 5, 6],
       [0, 0, 9]], dtype=int32)

When k > 0:

>>> jnp.triu_indices_from(arr, k=1)
(Array([0, 0, 1], dtype=int32), Array([1, 2, 2], dtype=int32))

When k < 0:

>>> jnp.triu_indices_from(arr, k=-1)
(Array([0, 0, 0, 1, 1, 1, 2, 2], dtype=int32), Array([0, 1, 2, 0, 1, 2, 1, 2], dtype=int32))
Parameters:
Return type:

tuple[Array, Array]

quchip.declarative.qnp.true_divide(x1, x2, /)

Calculates the division of x1 by x2 element-wise

JAX implementation of numpy.true_divide().

Args:

x1: Input array, the dividend x2: Input array, the divisor

Returns:

An array containing the elementwise quotients, will always use floating point division.

Examples:
>>> x1 = jnp.array([3, 4, 5])
>>> x2 = 2
>>> jnp.true_divide(x1, x2)
Array([1.5, 2. , 2.5], dtype=float32)

>>> x1 = 24
>>> x2 = jnp.array([3, 4, 6j])
>>> jnp.true_divide(x1, x2)
Array([8.+0.j, 6.+0.j, 0.-4.j], dtype=complex64)

>>> x1 = jnp.array([1j, 9+5j, -4+2j])
>>> x2 = 3j
>>> jnp.true_divide(x1, x2)
Array([0.33333334+0.j       , 1.6666666 -3.j       ,
       0.6666667 +1.3333334j], dtype=complex64)
See Also:

jax.numpy.floor_divide() for integer division

Parameters:
Return type:

Array

quchip.declarative.qnp.trunc(x)

Round input to the nearest integer towards zero.

JAX implementation of numpy.trunc().

Args:

x: input array or scalar.

Returns:

An array with same shape and dtype as x containing the rounded values.

See also:

  • jax.numpy.fix(): Rounds the input to the nearest integer towards zero.

  • jax.numpy.ceil(): Rounds the input up to the nearest integer.

  • jax.numpy.floor(): Rounds the input down to the nearest integer.

Examples:
>>> key = jax.random.key(42)
>>> x = jax.random.uniform(key, (3, 3), minval=-10, maxval=10)
>>> with jnp.printoptions(precision=2, suppress=True):
...     print(x)
[[-0.23  3.6   2.33]
 [ 1.22 -0.99  1.72]
 [-8.5   5.5   3.98]]
>>> jnp.trunc(x)
Array([[-0.,  3.,  2.],
       [ 1., -0.,  1.],
       [-8.,  5.,  3.]], dtype=float32)
Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex)

Return type:

Array

class quchip.declarative.qnp.ufunc(func, /, nin, nout, *, name=None, nargs=None, identity=None, call=None, reduce=None, accumulate=None, at=None, reduceat=None)

Bases: object

Universal functions which operation element-by-element on arrays.

JAX implementation of numpy.ufunc.

This is a class for JAX-backed implementations of NumPy’s ufunc APIs. Most users will never need to instantiate ufunc, but rather will use the pre-defined ufuncs in jax.numpy.

For constructing your own ufuncs, see jax.numpy.frompyfunc().

Examples:

Universal functions are functions that apply element-wise to broadcasted arrays, but they also come with a number of extra attributes and methods.

As an example, consider the function jax.numpy.add. The object acts as a function that applies addition to broadcasted arrays in an element-wise manner:

>>> x = jnp.array([1, 2, 3, 4, 5])
>>> jnp.add(x, 1)
Array([2, 3, 4, 5, 6], dtype=int32)

Each ufunc object includes a number of attributes that describe its behavior:

>>> jnp.add.nin  # number of inputs
2
>>> jnp.add.nout  # number of outputs
1
>>> jnp.add.identity  # identity value, or None if no identity exists
0

Binary ufuncs like jax.numpy.add include number of methods to apply the function to arrays in different manners.

The outer() method applies the function to the pair-wise outer-product of the input array values:

>>> jnp.add.outer(x, x)
Array([[ 2,  3,  4,  5,  6],
       [ 3,  4,  5,  6,  7],
       [ 4,  5,  6,  7,  8],
       [ 5,  6,  7,  8,  9],
       [ 6,  7,  8,  9, 10]], dtype=int32)

The ufunc.reduce() method performs a reduction over the array. For example, jnp.add.reduce() is equivalent to jnp.sum:

>>> jnp.add.reduce(x)
Array(15, dtype=int32)

The ufunc.accumulate() method performs a cumulative reduction over the array. For example, jnp.add.accumulate() is equivalent to jax.numpy.cumulative_sum():

>>> jnp.add.accumulate(x)
Array([ 1,  3,  6, 10, 15], dtype=int32)

The ufunc.at() method applies the function at particular indices in the array; for jnp.add the computation is similar to jax.lax.scatter_add():

>>> jnp.add.at(x, 0, 100, inplace=False)
Array([101,   2,   3,   4,   5], dtype=int32)

And the ufunc.reduceat() method performs a number of reduce operations between specified indices of an array; for jnp.add the operation is similar to jax.ops.segment_sum():

>>> jnp.add.reduceat(x, jnp.array([0, 2]))
Array([ 3, 12], dtype=int32)

In this case, the first element is x[0:2].sum(), and the second element is x[2:].sum().

Parameters:

  • func (Callable[..., Any])

  • nin (int)

  • nout (int)

  • name (str | None)

  • nargs (int | None)

  • identity (Any)

  • call (Callable[..., Any] | None)

  • reduce (Callable[..., Any] | None)

  • accumulate (Callable[..., Any] | None)

  • at (Callable[..., Any] | None)

  • reduceat (Callable[..., Any] | None)

accumulate(a, axis=0, dtype=None, out=None)[source]

Accumulate operation derived from binary ufunc.

JAX implementation of numpy.ufunc.accumulate().

Args:

a: N-dimensional array over which to accumulate. axis: integer axis over which accumulation will be performed (default = 0) dtype: optionally specify the type of the output array. out: Unused by JAX

Returns:

An array containing the accumulated result.

Examples:

Consider the following array:

>>> x = jnp.array([[1, 2, 3],
...                [4, 5, 6]])

jax.numpy.add.accumulate() is equivalent to jax.numpy.cumsum() along the specified axis: >>> jnp.add.accumulate(x, axis=1) Array([[ 1, 3, 6],

[ 4, 9, 15]], dtype=int32)

>>> jnp.cumsum(x, axis=1)
Array([[ 1,  3,  6],
       [ 4,  9, 15]], dtype=int32)

Similarly, jax.numpy.multiply.accumulate() is equivalent to jax.numpy.cumprod() along the specified axis:

>>> jnp.multiply.accumulate(x, axis=1)
Array([[  1,   2,   6],
       [  4,  20, 120]], dtype=int32)
>>> jnp.cumprod(x, axis=1)
Array([[  1,   2,   6],
       [  4,  20, 120]], dtype=int32)

For other binary ufuncs, the accumulation is an operation not available via standard APIs. For example, jax.numpy.bitwise_or.accumulate() is essentially a bitwise cumulative any:

>>> jnp.bitwise_or.accumulate(x, axis=1)
Array([[1, 3, 3],
       [4, 5, 7]], dtype=int32)
Parameters:
Return type:

Array

at(a, indices, b=None, /, *, inplace=True)[source]

Update elements of an array via the specified unary or binary ufunc.

JAX implementation of numpy.ufunc.at().

Note:

numpy.ufunc.at() mutates arrays in-place. JAX arrays are immutable, so jax.numpy.ufunc.at() cannot replicate these semantics. Instead, JAX will return the updated value, but requires explicitly passing inplace=False as a reminder of this difference.

Args:

a: N-dimensional array to update indices: index, slice, or tuple of indices and slices. b: array of values for binary ufunc updates. inplace: must be set to False to indicate that an updated copy will be returned.

Returns:

an updated copy of the input array.

Examples:

Add numbers to specified indices:

>>> x = jnp.ones(10, dtype=int)
>>> indices = jnp.array([2, 5, 7])
>>> values = jnp.array([10, 20, 30])
>>> jnp.add.at(x, indices, values, inplace=False)
Array([ 1,  1, 11,  1,  1, 21,  1, 31,  1,  1], dtype=int32)

This is roughly equivalent to JAX’s jax.numpy.ndarray.at() method called this way:

>>> x.at[indices].add(values)
Array([ 1,  1, 11,  1,  1, 21,  1, 31,  1,  1], dtype=int32)
Parameters:
Return type:

Array

property identity
property nargs
property nin
property nout
outer(A, B, /)[source]

Apply the function to all pairs of values in A and B.

JAX implementation of numpy.ufunc.outer().

Args:

A: N-dimensional array B: N-dimensional array

Returns:

An array of shape tuple(*A.shape, *B.shape)

Examples:

A times-table for integers 1…10 created via jax.numpy.multiply.outer():

>>> x = jnp.arange(1, 11)
>>> print(jnp.multiply.outer(x, x))
[[  1   2   3   4   5   6   7   8   9  10]
 [  2   4   6   8  10  12  14  16  18  20]
 [  3   6   9  12  15  18  21  24  27  30]
 [  4   8  12  16  20  24  28  32  36  40]
 [  5  10  15  20  25  30  35  40  45  50]
 [  6  12  18  24  30  36  42  48  54  60]
 [  7  14  21  28  35  42  49  56  63  70]
 [  8  16  24  32  40  48  56  64  72  80]
 [  9  18  27  36  45  54  63  72  81  90]
 [ 10  20  30  40  50  60  70  80  90 100]]

For input arrays with N and M dimensions respectively, the output will have dimension N + M:

>>> x = jnp.ones((1, 3, 5))
>>> y = jnp.ones((2, 4))
>>> jnp.add.outer(x, y).shape
(1, 3, 5, 2, 4)
Parameters:
Return type:

Array

reduce(a, axis=0, dtype=None, out=None, keepdims=False, initial=None, where=None)[source]

Reduction operation derived from a binary function.

JAX implementation of numpy.ufunc.reduce().

Args:

a: Input array. axis: integer specifying the axis over which to reduce. default=0 dtype: optionally specify the type of the output array. out: Unused by JAX keepdims: If True, reduced axes are left in the result with size 1.

If False (default) then reduced axes are squeezed out.

initial: int or array, Default=None. Initial value for the reduction. where: boolean mask, default=None. The elements to be used in the sum. Array

should be broadcast compatible to the input.

Returns:

array containing the result of the reduction operation.

Examples:

Consider the following array:

>>> x = jnp.array([[1, 2, 3],
...                [4, 5, 6]])

jax.numpy.add.reduce() is equivalent to jax.numpy.sum() along axis=0:

>>> jnp.add.reduce(x)
Array([5, 7, 9], dtype=int32)
>>> x.sum(0)
Array([5, 7, 9], dtype=int32)

Similarly, jax.numpy.logical_and.reduce() is equivalent to jax.numpy.all():

>>> jnp.logical_and.reduce(x > 2)
Array([False, False,  True], dtype=bool)
>>> jnp.all(x > 2, axis=0)
Array([False, False,  True], dtype=bool)

Some reductions do not correspond to any built-in aggregation function; for example here is the reduction of jax.numpy.bitwise_or() along the first axis of x:

>>> jnp.bitwise_or.reduce(x, axis=1)
Array([3, 7], dtype=int32)
Parameters:
Return type:

Array

reduceat(a, indices, axis=0, dtype=None, out=None)[source]

Reduce an array between specified indices via a binary ufunc.

JAX implementation of numpy.ufunc.reduceat()

Args:

a: N-dimensional array to reduce indices: a 1-dimensional array of increasing integer values which encodes

segments of the array to be reduced.

axis: integer specifying the axis along which to reduce: default=0. dtype: optionally specify the dtype of the output array. out: unused by JAX

Returns:

An array containing the reduced values.

Examples:

The reduce method lets you efficiently compute reduction operations over array segments. For example:

>>> x = jnp.array([1, 2, 3, 4, 5, 6, 7, 8])
>>> indices = jnp.array([0, 2, 5])
>>> jnp.add.reduce(x, indices)
Array([ 3, 12, 21], dtype=int32)

This is more-or-less equivalent to the following:

>>> jnp.array([x[0:2].sum(), x[2:5].sum(), x[5:].sum()])
Array([ 3, 12, 21], dtype=int32)

For some binary ufuncs, JAX provides similar APIs within jax.ops. For example, jax.add.reduceat() is similar to jax.ops.segment_sum(), although in this case the segments are defined via an array of segment ids:

>>> segments = jnp.array([0, 0, 1, 1, 1, 2, 2, 2])
>>> jax.ops.segment_sum(x, segments)
Array([ 3, 12, 21], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.uint

alias of uint64

class quchip.declarative.qnp.uint16(x)

Bases: object

A JAX scalar constructor of type uint16.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype('uint16')
class quchip.declarative.qnp.uint2(x)

Bases: object

A JAX scalar constructor of type uint2.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype(uint2)
class quchip.declarative.qnp.uint32(x)

Bases: object

A JAX scalar constructor of type uint32.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype('uint32')
class quchip.declarative.qnp.uint4(x)

Bases: object

A JAX scalar constructor of type uint4.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype(uint4)
class quchip.declarative.qnp.uint64(x)

Bases: object

A JAX scalar constructor of type uint64.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype('uint64')
class quchip.declarative.qnp.uint8(x)

Bases: object

A JAX scalar constructor of type uint8.

While NumPy defines scalar types for each data type, JAX represents scalars as zero-dimensional arrays.

Parameters:

x (Any)

Return type:

Array

dtype = dtype('uint8')
quchip.declarative.qnp.union1d(ar1, ar2, *, size=None, fill_value=None)

Compute the set union of two 1D arrays.

JAX implementation of numpy.union1d().

Because the size of the output of union1d is data-dependent, the function is not typically compatible with jit() and other JAX transformations. The JAX version adds the optional size argument which must be specified statically for jnp.union1d to be used in such contexts.

Args:

ar1: first array of elements to be unioned. ar2: second array of elements to be unioned size: if specified, return only the first size sorted elements. If there are fewer

elements than size indicates, the return value will be padded with fill_value.

fill_value: when size is specified and there are fewer than the indicated number of

elements, fill the remaining entries fill_value. Defaults to the minimum value.

Returns:

an array containing the union of elements in the input array.

See also:

  • jax.numpy.intersect1d(): the set intersection of two 1D arrays.

  • jax.numpy.setxor1d(): the set XOR of two 1D arrays.

  • jax.numpy.setdiff1d(): the set difference of two 1D arrays.

Examples:

Computing the union of two arrays:

>>> ar1 = jnp.array([1, 2, 3, 4])
>>> ar2 = jnp.array([3, 4, 5, 6])
>>> jnp.union1d(ar1, ar2)
Array([1, 2, 3, 4, 5, 6], dtype=int32)

Because the output shape is dynamic, this will fail under jit() and other transformations:

>>> jax.jit(jnp.union1d)(ar1, ar2)
Traceback (most recent call last):
   ...
ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: traced array with shape int32[4].
The error occurred while tracing the function union1d at /Users/vanderplas/github/jax-ml/jax/jax/_src/numpy/setops.py:101 for jit. This concrete value was not available in Python because it depends on the value of the argument ar1.

In order to ensure statically-known output shapes, you can pass a static size argument:

>>> jit_union1d = jax.jit(jnp.union1d, static_argnames=['size'])
>>> jit_union1d(ar1, ar2, size=6)
Array([1, 2, 3, 4, 5, 6], dtype=int32)

If size is too small, the union is truncated:

>>> jit_union1d(ar1, ar2, size=4)
Array([1, 2, 3, 4], dtype=int32)

If size is too large, then the output is padded with fill_value:

>>> jit_union1d(ar1, ar2, size=8, fill_value=0)
Array([1, 2, 3, 4, 5, 6, 0, 0], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.unique(ar, return_index=False, return_inverse=False, return_counts=False, axis=None, *, equal_nan=True, size=None, fill_value=None, sorted=True)

Return the unique values from an array.

JAX implementation of numpy.unique().

Because the size of the output of unique is data-dependent, the function is not typically compatible with jit() and other JAX transformations. The JAX version adds the optional size argument which must be specified statically for jnp.unique to be used in such contexts.

Args:

ar: N-dimensional array from which unique values will be extracted. return_index: if True, also return the indices in ar where each value occurs return_inverse: if True, also return the indices that can be used to reconstruct

ar from the unique values.

return_counts: if True, also return the number of occurrences of each unique value. axis: if specified, compute unique values along the specified axis. If None (default),

then flatten ar before computing the unique values.

equal_nan: if True, consider NaN values equivalent when determining uniqueness. size: if specified, return only the first size sorted unique elements. If there are fewer

unique elements than size indicates, the return value will be padded with fill_value.

fill_value: when size is specified and there are fewer than the indicated number of

elements, fill the remaining entries fill_value. Defaults to the minimum unique value.

sorted: unused by JAX.

Returns:

An array or tuple of arrays, depending on the values of return_index, return_inverse, and return_counts. Returned values are

  • unique_values:

    if axis is None, a 1D array of length n_unique, If axis is specified, shape is (*ar.shape[:axis], n_unique, *ar.shape[axis + 1:]).

  • unique_index:

    (returned only if return_index is True) An array of shape (n_unique,). Contains the indices of the first occurrence of each unique value in ar. For 1D inputs, ar[unique_index] is equivalent to unique_values.

  • unique_inverse:

    (returned only if return_inverse is True) An array of shape (ar.size,) if axis is None, or of shape (ar.shape[axis],) if axis is specified. Contains the indices within unique_values of each value in ar. For 1D inputs, unique_values[unique_inverse] is equivalent to ar.

  • unique_counts:

    (returned only if return_counts is True) An array of shape (n_unique,). Contains the number of occurrences of each unique value in ar.

See also:

  • jax.numpy.unique_counts(): shortcut to unique(arr, return_counts=True).

  • jax.numpy.unique_inverse(): shortcut to unique(arr, return_inverse=True).

  • jax.numpy.unique_all(): shortcut to unique with all return values.

  • jax.numpy.unique_values(): like unique, but no optional return values.

Examples:
>>> x = jnp.array([3, 4, 1, 3, 1])
>>> jnp.unique(x)
Array([1, 3, 4], dtype=int32)

JIT compilation & the size argument

If you try this under jit() or another transformation, you will get an error because the output shape is dynamic:

>>> jax.jit(jnp.unique)(x)
Traceback (most recent call last):
   ...
jax.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: traced array with shape int32[5].
The error arose for the first argument of jnp.unique(). To make jnp.unique() compatible with JIT and other transforms, you can specify a concrete value for the size argument, which will determine the output size.

The issue is that the output of transformed functions must have static shapes. In order to make this work, you can pass a static size parameter:

>>> jit_unique = jax.jit(jnp.unique, static_argnames=['size'])
>>> jit_unique(x, size=3)
Array([1, 3, 4], dtype=int32)

If your static size is smaller than the true number of unique values, they will be truncated.

>>> jit_unique(x, size=2)
Array([1, 3], dtype=int32)

If the static size is larger than the true number of unique values, they will be padded with fill_value, which defaults to the minimum unique value:

>>> jit_unique(x, size=5)
Array([1, 3, 4, 1, 1], dtype=int32)
>>> jit_unique(x, size=5, fill_value=0)
Array([1, 3, 4, 0, 0], dtype=int32)

Multi-dimensional unique values

If you pass a multi-dimensional array to unique, it will be flattened by default:

>>> M = jnp.array([[1, 2],
...                [2, 3],
...                [1, 2]])
>>> jnp.unique(M)
Array([1, 2, 3], dtype=int32)

If you pass an axis keyword, you can find unique slices of the array along that axis:

>>> jnp.unique(M, axis=0)
Array([[1, 2],
       [2, 3]], dtype=int32)

Returning indices

If you set return_index=True, then unique returns the indices of the first occurrence of each unique value:

>>> x = jnp.array([3, 4, 1, 3, 1])
>>> values, indices = jnp.unique(x, return_index=True)
>>> print(values)
[1 3 4]
>>> print(indices)
[2 0 1]
>>> jnp.all(values == x[indices])
Array(True, dtype=bool)

In multiple dimensions, the unique values can be extracted with jax.numpy.take() evaluated along the specified axis:

>>> values, indices = jnp.unique(M, axis=0, return_index=True)
>>> jnp.all(values == jnp.take(M, indices, axis=0))
Array(True, dtype=bool)

Returning inverse

If you set return_inverse=True, then unique returns the indices within the unique values for every entry in the input array:

>>> x = jnp.array([3, 4, 1, 3, 1])
>>> values, inverse = jnp.unique(x, return_inverse=True)
>>> print(values)
[1 3 4]
>>> print(inverse)
[1 2 0 1 0]
>>> jnp.all(values[inverse] == x)
Array(True, dtype=bool)

In multiple dimensions, the input can be reconstructed using jax.numpy.take():

>>> values, inverse = jnp.unique(M, axis=0, return_inverse=True)
>>> jnp.all(jnp.take(values, inverse, axis=0) == M)
Array(True, dtype=bool)

Returning counts

If you set return_counts=True, then unique returns the number of occurrences within the input for every unique value:

>>> x = jnp.array([3, 4, 1, 3, 1])
>>> values, counts = jnp.unique(x, return_counts=True)
>>> print(values)
[1 3 4]
>>> print(counts)
[2 2 1]

For multi-dimensional arrays, this also returns a 1D array of counts indicating number of occurrences along the specified axis:

>>> values, counts = jnp.unique(M, axis=0, return_counts=True)
>>> print(values)
[[1 2]
 [2 3]]
>>> print(counts)
[2 1]
Parameters:
quchip.declarative.qnp.unique_all(x, /, *, size=None, fill_value=None)

Return unique values from x, along with indices, inverse indices, and counts.

JAX implementation of numpy.unique_all(); this is equivalent to calling jax.numpy.unique() with return_index, return_inverse, return_counts, and equal_nan set to True.

Because the size of the output of unique_all is data-dependent, the function is not typically compatible with jit() and other JAX transformations. The JAX version adds the optional size argument which must be specified statically for jnp.unique to be used in such contexts.

Args:

x: N-dimensional array from which unique values will be extracted. size: if specified, return only the first size sorted unique elements. If there are fewer

unique elements than size indicates, the return value will be padded with fill_value.

fill_value: when size is specified and there are fewer than the indicated number of

elements, fill the remaining entries fill_value. Defaults to the minimum unique value.

Returns:

A tuple (values, indices, inverse_indices, counts), with the following properties:

  • values:

    an array of shape (n_unique,) containing the unique values from x.

  • indices:

    An array of shape (n_unique,). Contains the indices of the first occurrence of each unique value in x. For 1D inputs, x[indices] is equivalent to values.

  • inverse_indices:

    An array of shape x.shape. Contains the indices within values of each value in x. For 1D inputs, values[inverse_indices] is equivalent to x.

  • counts:

    An array of shape (n_unique,). Contains the number of occurrences of each unique value in x.

See also:

  • jax.numpy.unique(): general function for computing unique values.

  • jax.numpy.unique_values(): compute only values.

  • jax.numpy.unique_counts(): compute only values and counts.

  • jax.numpy.unique_inverse(): compute only values and inverse.

Examples:

Here we compute the unique values in a 1D array:

>>> x = jnp.array([3, 4, 1, 3, 1])
>>> result = jnp.unique_all(x)

The result is a NamedTuple with four named attributes. The values attribute contains the unique values from the array:

>>> result.values
Array([1, 3, 4], dtype=int32)

The indices attribute contains the indices of the unique values within the input array:

>>> result.indices
Array([2, 0, 1], dtype=int32)
>>> jnp.all(result.values == x[result.indices])
Array(True, dtype=bool)

The inverse_indices attribute contains the indices of the input within values:

>>> result.inverse_indices
Array([1, 2, 0, 1, 0], dtype=int32)
>>> jnp.all(x == result.values[result.inverse_indices])
Array(True, dtype=bool)

The counts attribute contains the counts of each unique value in the input:

>>> result.counts
Array([2, 2, 1], dtype=int32)

For examples of the size and fill_value arguments, see jax.numpy.unique().

Parameters:
Return type:

_UniqueAllResult

quchip.declarative.qnp.unique_counts(x, /, *, size=None, fill_value=None)

Return unique values from x, along with counts.

JAX implementation of numpy.unique_counts(); this is equivalent to calling jax.numpy.unique() with return_counts and equal_nan set to True.

Because the size of the output of unique_counts is data-dependent, the function is not typically compatible with jit() and other JAX transformations. The JAX version adds the optional size argument which must be specified statically for jnp.unique to be used in such contexts.

Args:

x: N-dimensional array from which unique values will be extracted. size: if specified, return only the first size sorted unique elements. If there are fewer

unique elements than size indicates, the return value will be padded with fill_value.

fill_value: when size is specified and there are fewer than the indicated number of

elements, fill the remaining entries fill_value. Defaults to the minimum unique value.

Returns:

A tuple (values, counts), with the following properties:

  • values:

    an array of shape (n_unique,) containing the unique values from x.

  • counts:

    An array of shape (n_unique,). Contains the number of occurrences of each unique value in x.

See also:

  • jax.numpy.unique(): general function for computing unique values.

  • jax.numpy.unique_values(): compute only values.

  • jax.numpy.unique_inverse(): compute only values and inverse.

  • jax.numpy.unique_all(): compute values, indices, inverse_indices, and counts.

Examples:

Here we compute the unique values in a 1D array:

>>> x = jnp.array([3, 4, 1, 3, 1])
>>> result = jnp.unique_counts(x)

The result is a NamedTuple with two named attributes. The values attribute contains the unique values from the array:

>>> result.values
Array([1, 3, 4], dtype=int32)

The counts attribute contains the counts of each unique value in the input:

>>> result.counts
Array([2, 2, 1], dtype=int32)

For examples of the size and fill_value arguments, see jax.numpy.unique().

Parameters:
Return type:

_UniqueCountsResult

quchip.declarative.qnp.unique_inverse(x, /, *, size=None, fill_value=None)

Return unique values from x, along with indices, inverse indices, and counts.

JAX implementation of numpy.unique_inverse(); this is equivalent to calling jax.numpy.unique() with return_inverse and equal_nan set to True.

Because the size of the output of unique_inverse is data-dependent, the function is not typically compatible with jit() and other JAX transformations. The JAX version adds the optional size argument which must be specified statically for jnp.unique to be used in such contexts.

Args:

x: N-dimensional array from which unique values will be extracted. size: if specified, return only the first size sorted unique elements. If there are fewer

unique elements than size indicates, the return value will be padded with fill_value.

fill_value: when size is specified and there are fewer than the indicated number of

elements, fill the remaining entries fill_value. Defaults to the minimum unique value.

Returns:

A tuple (values, indices, inverse_indices, counts), with the following properties:

  • values:

    an array of shape (n_unique,) containing the unique values from x.

  • inverse_indices:

    An array of shape x.shape. Contains the indices within values of each value in x. For 1D inputs, values[inverse_indices] is equivalent to x.

See also:

  • jax.numpy.unique(): general function for computing unique values.

  • jax.numpy.unique_values(): compute only values.

  • jax.numpy.unique_counts(): compute only values and counts.

  • jax.numpy.unique_all(): compute values, indices, inverse_indices, and counts.

Examples:

Here we compute the unique values in a 1D array:

>>> x = jnp.array([3, 4, 1, 3, 1])
>>> result = jnp.unique_inverse(x)

The result is a NamedTuple with two named attributes. The values attribute contains the unique values from the array:

>>> result.values
Array([1, 3, 4], dtype=int32)

The indices attribute contains the indices of the unique values within the input array:

The inverse_indices attribute contains the indices of the input within values:

>>> result.inverse_indices
Array([1, 2, 0, 1, 0], dtype=int32)
>>> jnp.all(x == result.values[result.inverse_indices])
Array(True, dtype=bool)

For examples of the size and fill_value arguments, see jax.numpy.unique().

Parameters:
Return type:

_UniqueInverseResult

quchip.declarative.qnp.unique_values(x, /, *, size=None, fill_value=None)

Return unique values from x, along with indices, inverse indices, and counts.

JAX implementation of numpy.unique_values(); this is equivalent to calling jax.numpy.unique() with equal_nan set to True.

Because the size of the output of unique_values is data-dependent, the function is not typically compatible with jit() and other JAX transformations. The JAX version adds the optional size argument which must be specified statically for jnp.unique to be used in such contexts.

Args:

x: N-dimensional array from which unique values will be extracted. size: if specified, return only the first size sorted unique elements. If there are fewer

unique elements than size indicates, the return value will be padded with fill_value.

fill_value: when size is specified and there are fewer than the indicated number of

elements, fill the remaining entries fill_value. Defaults to the minimum unique value.

Returns:

An array values of shape (n_unique,) containing the unique values from x.

See also:

  • jax.numpy.unique(): general function for computing unique values.

  • jax.numpy.unique_values(): compute only values.

  • jax.numpy.unique_counts(): compute only values and counts.

  • jax.numpy.unique_inverse(): compute only values and inverse.

Examples:

Here we compute the unique values in a 1D array:

>>> x = jnp.array([3, 4, 1, 3, 1])
>>> jnp.unique_values(x)
Array([1, 3, 4], dtype=int32)

For examples of the size and fill_value arguments, see jax.numpy.unique().

Parameters:
Return type:

Array

quchip.declarative.qnp.unpackbits(a, axis=None, count=None, bitorder='big')

Unpack the bits in a uint8 array.

JAX implementation of numpy.unpackbits().

Args:

a: N-dimensional array of type uint8. axis: optional axis along which to unpack. If not specified, a will

be flattened

count: specify the number of bits to unpack (if positive) or the number

of bits to trim from the end (if negative).

bitorder: "big" (default) or "little": specify whether the bit order

is big-endian or little-endian.

Returns:

a uint8 array of unpacked bits.

See also:

  • jax.numpy.packbits(): this inverse of unpackbits.

Examples:

Unpacking bits from a scalar:

>>> jnp.unpackbits(jnp.uint8(27))  # big-endian by default
Array([0, 0, 0, 1, 1, 0, 1, 1], dtype=uint8)
>>> jnp.unpackbits(jnp.uint8(27), bitorder="little")
Array([1, 1, 0, 1, 1, 0, 0, 0], dtype=uint8)

Compare this to the Python binary representation:

>>> 0b00011011
27

Unpacking bits along an axis:

>>> vals = jnp.array([[154],
...                   [ 49]], dtype='uint8')
>>> bits = jnp.unpackbits(vals, axis=1)
>>> bits
Array([[1, 0, 0, 1, 1, 0, 1, 0],
       [0, 0, 1, 1, 0, 0, 0, 1]], dtype=uint8)

Using packbits() to invert this:

>>> jnp.packbits(bits, axis=1)
Array([[154],
       [ 49]], dtype=uint8)

The count keyword lets unpackbits serve as an inverse of packbits in cases where not all bits are present:

>>> bits = jnp.array([1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1])  # 11 bits
>>> vals = jnp.packbits(bits)
>>> vals
Array([219,  96], dtype=uint8)
>>> jnp.unpackbits(vals)  # 16 zero-padded bits
Array([1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 0, 0, 0, 0, 0], dtype=uint8)
>>> jnp.unpackbits(vals, count=11)  # specify 11 output bits
Array([1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1], dtype=uint8)
>>> jnp.unpackbits(vals, count=-5)  # specify 5 bits to be trimmed
Array([1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1], dtype=uint8)
Parameters:
Return type:

Array

quchip.declarative.qnp.unravel_index(indices, shape)

Convert flat indices into multi-dimensional indices.

JAX implementation of numpy.unravel_index(). The JAX version differs in its treatment of out-of-bound indices: unlike NumPy, negative indices are supported, and out-of-bound indices are clipped to the nearest valid value.

Args:

indices: integer array of flat indices shape: shape of multidimensional array to index into

Returns:

Tuple of unraveled indices

See also:

jax.numpy.ravel_multi_index(): Inverse of this function.

Examples:

Start with a 1D array values and indices:

>>> x = jnp.array([2., 3., 4., 5., 6., 7.])
>>> indices = jnp.array([1, 3, 5])
>>> print(x[indices])
[3. 5. 7.]

Now if x is reshaped, unravel_indices can be used to convert the flat indices into a tuple of indices that access the same entries:

>>> shape = (2, 3)
>>> x_2D = x.reshape(shape)
>>> indices_2D = jnp.unravel_index(indices, shape)
>>> indices_2D
(Array([0, 1, 1], dtype=int32), Array([1, 0, 2], dtype=int32))
>>> print(x_2D[indices_2D])
[3. 5. 7.]

The inverse function, ravel_multi_index, can be used to obtain the original indices:

>>> jnp.ravel_multi_index(indices_2D, shape)
Array([1, 3, 5], dtype=int32)
Parameters:
Return type:

tuple[Array, …]

class quchip.declarative.qnp.unsignedinteger

Bases: integer

Abstract base class of all unsigned integer scalar types.

quchip.declarative.qnp.unstack(x, /, *, axis=0)

Unstack an array along an axis.

JAX implementation of array_api.unstack().

Args:

x: array to unstack. Must have x.ndim >= 1. axis: integer axis along which to unstack. Must satisfy

-x.ndim <= axis < x.ndim.

Returns:

tuple of unstacked arrays.

See also:

  • jax.numpy.stack(): inverse of unstack

  • jax.numpy.split(): split array into batches along an axis.

Examples:
>>> arr = jnp.array([[1, 2, 3],
...                  [4, 5, 6]])
>>> arrs = jnp.unstack(arr)
>>> print(*arrs)
[1 2 3] [4 5 6]

stack() provides the inverse of this:

>>> jnp.stack(arrs)
Array([[1, 2, 3],
       [4, 5, 6]], dtype=int32)
Parameters:
Return type:

tuple[Array, …]

quchip.declarative.qnp.unwrap(p, discont=None, axis=-1, period=6.283185307179586)

Unwrap a periodic signal.

JAX implementation of numpy.unwrap().

Args:

p: input array discont: the maximum allowable discontinuity in the sequence. The

default is period / 2

axis: the axis along which to unwrap; defaults to -1 period: the period of the signal, which defaults to \(2\pi\)

Returns:

An unwrapped copy of p.

Examples:

Consider a situation in which you are making measurements of the position of a rotating disk via the x and y locations of some point on that disk. The underlying variable is an always-increating angle which we’ll generate this way, using degrees for ease of representation:

>>> rng = np.random.default_rng(0)
>>> theta = rng.integers(0, 90, size=(20,)).cumsum()
>>> theta
array([ 76, 133, 179, 203, 230, 233, 239, 240, 255, 328, 386, 468, 513,
       567, 654, 719, 775, 823, 873, 957])

Our observations of this angle are the x and y coordinates, given by the sine and cosine of this underlying angle:

>>> x, y = jnp.sin(jnp.deg2rad(theta)), jnp.cos(jnp.deg2rad(theta))

Now, say that given these x and y coordinates, we wish to recover the original angle theta. We might do this via the atan2() function:

>>> theta_out = jnp.rad2deg(jnp.atan2(x, y)).round()
>>> theta_out
Array([  76.,  133.,  179., -157., -130., -127., -121., -120., -105.,
        -32.,   26.,  108.,  153., -153.,  -66.,   -1.,   55.,  103.,
        153., -123.], dtype=float32)

The first few values match the input angle theta above, but after this the values are wrapped because the sin and cos observations obscure the phase information. The purpose of the unwrap() function is to recover the original signal from this wrapped view of it:

>>> jnp.unwrap(theta_out, period=360)
Array([ 76., 133., 179., 203., 230., 233., 239., 240., 255., 328., 386.,
       468., 513., 567., 654., 719., 775., 823., 873., 957.],      dtype=float32)

It does this by assuming that the true underlying sequence does not differ by more than discont (which defaults to period / 2) within a single step, and when it encounters a larger discontinuity it adds factors of the period to the data. For periodic signals that satisfy this assumption, unwrap() can recover the original phased signal.

Parameters:
Return type:

Array

quchip.declarative.qnp.vander(x, N=None, increasing=False)

Generate a Vandermonde matrix.

JAX implementation of numpy.vander().

Args:

x: input array. Must have x.ndim == 1. N: int, optional, default=None. Specifies the number of the columns the

output matrix. If not specified, N = len(x).

increasing: bool, optional, default=False. Specifies the order of the powers

of the columns. If True, the powers increase from left to right, \([x^0, x^1, ..., x^{(N-1)}]\). By default, the powers decrease from left to right \([x^{(N-1)}, ..., x^1, x^0]\).

Returns:

An array of shape [len(x), N] containing the generated Vandermonde matrix.

Examples:
>>> x = jnp.array([1, 2, 3, 4])
>>> jnp.vander(x)
Array([[ 1,  1,  1,  1],
       [ 8,  4,  2,  1],
       [27,  9,  3,  1],
       [64, 16,  4,  1]], dtype=int32)

If N = 2, generates a Vandermonde matrix with 2 columns.

>>> jnp.vander(x, N=2)
Array([[1, 1],
       [2, 1],
       [3, 1],
       [4, 1]], dtype=int32)

Generates the Vandermonde matrix in increasing order of powers, when increasing=True.

>>> jnp.vander(x, increasing=True)
Array([[ 1,  1,  1,  1],
       [ 1,  2,  4,  8],
       [ 1,  3,  9, 27],
       [ 1,  4, 16, 64]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=False, *, where=None, correction=None)

Compute the variance along a given axis.

JAX implementation of numpy.var().

Args:

a: input array. axis: optional, int or sequence of ints, default=None. Axis along which the

variance is computed. If None, variance is computed along all the axes.

dtype: The type of the output array. Default=None. ddof: int, default=0. Degrees of freedom. The divisor in the variance computation

is N-ddof, N is number of elements along given axis.

keepdims: bool, default=False. If true, reduced axes are left in the result

with size 1.

where: optional, boolean array, default=None. The elements to be used in the

variance. Array should be broadcast compatible to the input.

correction: int or float, default=None. Alternative name for ddof.

Both ddof and correction can’t be provided simultaneously.

out: Unused by JAX.

Returns:

An array of the variance along the given axis.

See also:

  • jax.numpy.mean(): Compute the mean of array elements over a given axis.

  • jax.numpy.std(): Compute the standard deviation of array elements over given axis.

  • jax.numpy.nanvar(): Compute the variance along a given axis, ignoring NaNs values.

  • jax.numpy.nanstd(): Computed the standard deviation of a given axis, ignoring NaN values.

Examples:

By default, jnp.var computes the variance along all axes.

>>> x = jnp.array([[1, 3, 4, 2],
...                [5, 2, 6, 3],
...                [8, 4, 2, 9]])
>>> with jnp.printoptions(precision=2, suppress=True):
...   jnp.var(x)
Array(5.74, dtype=float32)

If axis=1, variance is computed along axis 1.

>>> jnp.var(x, axis=1)
Array([1.25  , 2.5   , 8.1875], dtype=float32)

To preserve the dimensions of input, you can set keepdims=True.

>>> jnp.var(x, axis=1, keepdims=True)
Array([[1.25  ],
       [2.5   ],
       [8.1875]], dtype=float32)

If ddof=1:

>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jnp.var(x, axis=1, keepdims=True, ddof=1))
[[ 1.67]
 [ 3.33]
 [10.92]]

To include specific elements of the array to compute variance, you can use where.

>>> where = jnp.array([[1, 0, 1, 0],
...                    [0, 1, 1, 0],
...                    [1, 1, 1, 0]], dtype=bool)
>>> with jnp.printoptions(precision=2, suppress=True):
...   print(jnp.var(x, axis=1, keepdims=True, where=where))
[[2.25]
 [4.  ]
 [6.22]]
Parameters:
Return type:

Array

quchip.declarative.qnp.vdot(a, b, *, precision=None, preferred_element_type=None)

Perform a conjugate multiplication of two 1D vectors.

JAX implementation of numpy.vdot().

Args:

a: first input array, if not 1D it will be flattened. b: second input array, if not 1D it will be flattened. Must have a.size == b.size. precision: either None (default), which means the default precision for

the backend, a Precision enum value (Precision.DEFAULT, Precision.HIGH or Precision.HIGHEST) or a tuple of two such values indicating precision of a and b.

preferred_element_type: either None (default), which means the default

accumulation type for the input types, or a datatype, indicating to accumulate results to and return a result with that datatype.

Returns:

Scalar array (shape ()) containing the conjugate vector product of the inputs.

See Also:

  • jax.numpy.vecdot(): batched vector product.

  • jax.numpy.matmul(): general matrix multiplication.

  • jax.lax.dot_general(): general N-dimensional batched dot product.

Examples:
>>> x = jnp.array([1j, 2j, 3j])
>>> y = jnp.array([1., 2., 3.])
>>> jnp.vdot(x, y)
Array(0.-14.j, dtype=complex64)

Note the difference between this and dot(), which does not conjugate the first input when complex:

>>> jnp.dot(x, y)
Array(0.+14.j, dtype=complex64)
Parameters:
Return type:

Array

quchip.declarative.qnp.vecdot(x1, x2, /, *, axis=-1, precision=None, preferred_element_type=None)

Perform a conjugate multiplication of two batched vectors.

JAX implementation of numpy.vecdot().

Args:

a: left-hand side array. b: right-hand side array. Size of b[axis] must match size of a[axis],

and remaining dimensions must be broadcast-compatible.

axis: axis along which to compute the dot product (default: -1) precision: either None (default), which means the default precision for

the backend, a Precision enum value (Precision.DEFAULT, Precision.HIGH or Precision.HIGHEST) or a tuple of two such values indicating precision of a and b.

preferred_element_type: either None (default), which means the default

accumulation type for the input types, or a datatype, indicating to accumulate results to and return a result with that datatype.

Returns:

array containing the conjugate dot product of a and b along axis. The non-contracted dimensions are broadcast together.

See Also:

  • jax.numpy.vdot(): flattened vector product.

  • jax.numpy.vecmat(): vector-matrix product.

  • jax.numpy.matmul(): general matrix multiplication.

  • jax.lax.dot_general(): general N-dimensional batched dot product.

Examples:

Vector conjugate-dot product of two 1D arrays:

>>> a = jnp.array([1j, 2j, 3j])
>>> b = jnp.array([4., 5., 6.])
>>> jnp.linalg.vecdot(a, b)
Array(0.-32.j, dtype=complex64)

Batched vector dot product of two 2D arrays:

>>> a = jnp.array([[1, 2, 3],
...                [4, 5, 6]])
>>> b = jnp.array([[2, 3, 4]])
>>> jnp.linalg.vecdot(a, b, axis=-1)
Array([20, 47], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.vecmat(x1, x2, /)

Batched conjugate vector-matrix product.

JAX implementation of numpy.vecmat().

Args:

x1: array of shape (..., M). x2: array of shape (..., M, N). Leading dimensions must be broadcast-compatible

with leading dimensions of x1.

Returns:

An array of shape (..., N) containing the batched conjugate vector-matrix product.

See also:

  • jax.numpy.linalg.vecdot(): batched vector product.

  • jax.numpy.matvec(): matrix-vector product.

  • jax.numpy.matmul(): general matrix multiplication.

Examples:

Simple vector-matrix product:

>>> x1 = jnp.array([[1, 2, 3]])
>>> x2 = jnp.array([[4, 5],
...                 [6, 7],
...                 [8, 9]])
>>> jnp.vecmat(x1, x2)
Array([[40, 46]], dtype=int32)

Batched vector-matrix product:

>>> x1 = jnp.array([[1, 2, 3],
...                 [4, 5, 6]])
>>> jnp.vecmat(x1, x2)
Array([[ 40,  46],
       [ 94, 109]], dtype=int32)
Parameters:
Return type:

Array

quchip.declarative.qnp.vectorize(pyfunc, *, excluded=frozenset({}), signature=None)

Define a vectorized function with broadcasting.

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.

Args:

pyfunc: function to vectorize. excluded: optional set of integers representing positional arguments for

which the function will not be vectorized. These will be passed directly to pyfunc unmodified.

signature: optional generalized universal function signature, e.g.,

(m,n),(n)->(m) for vectorized matrix-vector multiplication. If provided, pyfunc will be called with (and expected to return) arrays with shapes given by the size of corresponding core dimensions. By default, pyfunc is assumed to take scalar arrays as input, and if signature is None, pyfunc can produce outputs of any shape.

Returns:

Vectorized version of the given function.

Examples:

Here are a few examples of how one could write vectorized linear algebra routines using vectorize():

>>> from functools import partial

>>> @partial(jnp.vectorize, signature='(k),(k)->(k)')
... def cross_product(a, b):
...   assert a.shape == b.shape and a.ndim == b.ndim == 1
...   return jnp.array([a[1] * b[2] - a[2] * b[1],
...                     a[2] * b[0] - a[0] * b[2],
...                     a[0] * b[1] - a[1] * b[0]])

>>> @partial(jnp.vectorize, signature='(n,m),(m)->(n)')
... def matrix_vector_product(matrix, vector):
...   assert matrix.ndim == 2 and matrix.shape[1:] == vector.shape
...   return matrix @ vector

These functions are only written to handle 1D or 2D arrays (the assert statements will never be violated), but with vectorize they support arbitrary dimensional inputs with NumPy style broadcasting, e.g.,

>>> cross_product(jnp.ones(3), jnp.ones(3)).shape
(3,)
>>> cross_product(jnp.ones((2, 3)), jnp.ones(3)).shape
(2, 3)
>>> cross_product(jnp.ones((1, 2, 3)), jnp.ones((2, 1, 3))).shape
(2, 2, 3)
>>> matrix_vector_product(jnp.ones(3), jnp.ones(3))
Traceback (most recent call last):
ValueError: input with shape (3,) does not have enough dimensions for all
core dimensions ('n', 'k') on vectorized function with excluded=frozenset()
and signature='(n,k),(k)->(k)'
>>> matrix_vector_product(jnp.ones((2, 3)), jnp.ones(3)).shape
(2,)
>>> matrix_vector_product(jnp.ones((2, 3)), jnp.ones((4, 3))).shape
(4, 2)

Note that this has different semantics than jnp.matmul:

>>> jnp.matmul(jnp.ones((2, 3)), jnp.ones((4, 3)))
Traceback (most recent call last):
TypeError: dot_general requires contracting dimensions to have the same shape, got [3] and [4].
quchip.declarative.qnp.vsplit(ary, indices_or_sections)

Split an array into sub-arrays vertically.

JAX implementation of numpy.vsplit().

Refer to the documentation of jax.numpy.split() for details; vsplit is equivalent to split with axis=0.

Examples:

1D array:

>>> x = jnp.array([1, 2, 3, 4, 5, 6])
>>> x1, x2 = jnp.vsplit(x, 2)
>>> print(x1, x2)
[1 2 3] [4 5 6]

2D array:

>>> x = jnp.array([[1, 2, 3, 4],
...                [5, 6, 7, 8]])
>>> x1, x2 = jnp.vsplit(x, 2)
>>> print(x1, x2)
[[1 2 3 4]] [[5 6 7 8]]
See also:

  • jax.numpy.split(): split an array along any axis.

  • jax.numpy.hsplit(): split horizontally, i.e. along axis=1

  • jax.numpy.dsplit(): split depth-wise, i.e. along axis=2

  • jax.numpy.array_split(): like split, but allows indices_or_sections to be an integer that does not evenly divide the size of the array.

Parameters:
Return type:

list[Array]

quchip.declarative.qnp.vstack(tup, dtype=None)

Vertically stack arrays.

JAX implementation of numpy.vstack().

For arrays of two or more dimensions, this is equivalent to jax.numpy.concatenate() with axis=0.

Args:
tup: a sequence of arrays to stack; each must have the same shape along all

but the first axis. If a single array is given it will be treated equivalently to tup = unstack(tup), but the implementation will avoid explicit unstacking.

dtype: optional dtype of the resulting array. If not specified, the dtype

will be determined via type promotion rules described in type-promotion.

Returns:

the stacked result.

See also:

  • jax.numpy.stack(): stack along arbitrary axes

  • jax.numpy.concatenate(): concatenation along existing axes.

  • jax.numpy.hstack(): stack horizontally, i.e. along axis 1.

  • jax.numpy.dstack(): stack depth-wise, i.e. along axis 2.

Examples:

Scalar values:

>>> jnp.vstack([1, 2, 3])
Array([[1],
       [2],
       [3]], dtype=int32, weak_type=True)

1D arrays:

>>> x = jnp.arange(4)
>>> y = jnp.ones(4)
>>> jnp.vstack([x, y])
Array([[0., 1., 2., 3.],
       [1., 1., 1., 1.]], dtype=float32)

2D arrays:

>>> x = x.reshape(1, 4)
>>> y = y.reshape(1, 4)
>>> jnp.vstack([x, y])
Array([[0., 1., 2., 3.],
       [1., 1., 1., 1.]], dtype=float32)
Parameters:
Return type:

Array

quchip.declarative.qnp.where(condition, x=None, y=None, /, *, size=None, fill_value=None)

Select elements from two arrays based on a condition.

JAX implementation of numpy.where().

Note

when only condition is provided, jnp.where(condition) is equivalent to jnp.nonzero(condition). For that case, refer to the documentation of jax.numpy.nonzero(). The docstring below focuses on the case where x and y are specified.

The three-term version of jnp.where lowers to jax.lax.select().

Args:
condition: boolean array. Must be broadcast-compatible with x and y when

they are specified.

x: arraylike. Should be broadcast-compatible with condition and y, and

typecast-compatible with y.

y: arraylike. Should be broadcast-compatible with condition and x, and

typecast-compatible with x.

size: integer, only referenced when x and y are None. For details,

see jax.numpy.nonzero().

fill_value: only referenced when x and y are None. For details,

see jax.numpy.nonzero().

Returns:

An array of dtype jnp.result_type(x, y) with values drawn from x where condition is True, and from y where condition is False. If x and y are None, the function behaves differently; see jax.numpy.nonzero() for a description of the return type.

See Also:

  • jax.numpy.nonzero()

  • jax.numpy.argwhere()

  • jax.lax.select()

Notes:

Special care is needed when the x or y input to jax.numpy.where() could have a value of NaN. Specifically, when a gradient is taken with jax.grad() (reverse-mode differentiation), a NaN in either x or y will propagate into the gradient, regardless of the value of condition. More information on this behavior and workarounds is available in the JAX FAQ.

Examples:

When x and y are not provided, where behaves equivalently to jax.numpy.nonzero():

>>> x = jnp.arange(10)
>>> jnp.where(x > 4)
(Array([5, 6, 7, 8, 9], dtype=int32),)
>>> jnp.nonzero(x > 4)
(Array([5, 6, 7, 8, 9], dtype=int32),)

When x and y are provided, where selects between them based on the specified condition:

>>> jnp.where(x > 4, x, 0)
Array([0, 0, 0, 0, 0, 5, 6, 7, 8, 9], dtype=int32)
quchip.declarative.qnp.zeros(shape, dtype=None, *, device=None)

Create an array full of zeros.

JAX implementation of numpy.zeros().

Args:

shape: int or sequence of ints specifying the shape of the created array. dtype: optional dtype for the created array; defaults to float32 or float64

depending on the X64 configuration (see default-dtypes).

device: (optional) Device or Sharding

to which the created array will be committed.

Returns:

Array of the specified shape and dtype, on the specified device if specified.

See also:

  • jax.numpy.zeros_like()

  • jax.numpy.empty()

  • jax.numpy.ones()

  • jax.numpy.full()

Examples:
>>> jnp.zeros(4)
Array([0., 0., 0., 0.], dtype=float32)
>>> jnp.zeros((2, 3), dtype=bool)
Array([[False, False, False],
       [False, False, False]], dtype=bool)
Parameters:

  • shape (Any)

  • dtype (str | type[Any] | dtype | SupportsDType | None)

  • device (Device | Sharding | None)

Return type:

Array

quchip.declarative.qnp.zeros_like(a, dtype=None, shape=None, *, device=None)

Create an array full of zeros with the same shape and dtype as an array.

JAX implementation of numpy.zeros_like().

Args:

a: Array-like object with shape and dtype attributes. shape: optionally override the shape of the created array. dtype: optionally override the dtype of the created array. device: (optional) Device or Sharding

to which the created array will be committed.

Returns:

Array of the specified shape and dtype, on the specified device if specified.

See also:

  • jax.numpy.zeros()

  • jax.numpy.empty_like()

  • jax.numpy.ones_like()

  • jax.numpy.full_like()

Examples:
>>> x = jnp.arange(4)
>>> jnp.zeros_like(x)
Array([0, 0, 0, 0], dtype=int32)
>>> jnp.zeros_like(x, dtype=bool)
Array([False, False, False, False], dtype=bool)
>>> jnp.zeros_like(x, shape=(2, 3))
Array([[0, 0, 0],
       [0, 0, 0]], dtype=int32)
Parameters:
Return type:

Array